From e87e66111808a55b85d5f8873a189a8f35de76a4 Mon Sep 17 00:00:00 2001 From: Mrunmayee Daware <165297479+Mrun25@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:45:51 -0700 Subject: [PATCH 01/36] fix(proxy): authorize the public name of a team-scoped model on batch retrieve/cancel (#47067) Backport of #47067 to stable/1.104.x. Cherry-picked from e979bd0ecf836976373e12624585962dd3356f3a (main). Conflict: main renamed the model access helper to check_model_access_helper, while stable/1.104.x still names it _check_model_access_helper, so the new owning-team check calls it by the line's name. Line adaptation: the line's test_auth_utils.py had no module-level Router import, which main's copy has and the picked helper uses, so it is added. --- litellm/proxy/auth/auth_checks.py | 30 +++++ .../proxy/auth/test_auth_utils.py | 119 +++++++++++++++++ .../proxy/batches_endpoints/test_endpoints.py | 123 ++++++++++++++++++ 3 files changed, 272 insertions(+) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 2ee64a71388..a52c2132bfa 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4352,6 +4352,24 @@ def _check_model_access_helper( return True +def _owning_team_names_for_team_deployment( + model: str, + llm_router: Router | None, + team_id: str | None, + team_model_aliases: Mapping[str, str] | None, +) -> tuple[str, ...]: + if llm_router is None or team_id is None or not model.startswith(f"model_name_{team_id}_"): + return () + deployment: Final = llm_router.get_deployment_by_model_group_name(model_group_name=model) + public_name: Final = ( + deployment.model_info.team_public_model_name + if deployment is not None and deployment.model_info.team_id == team_id + else None + ) + legacy_aliases: Final = tuple(alias for alias, target in (team_model_aliases or {}).items() if target == model) + return (*((public_name,) if public_name else ()), *legacy_aliases) + + def _can_object_call_model( model: str | list[str], llm_router: Router | None, @@ -4423,6 +4441,13 @@ def _can_object_call_model( *((global_or_router_alias_target,) if global_or_router_alias_target else ()), ) ) + owning_team_names: Final = ( + () + if key_alias_applied + else _owning_team_names_for_team_deployment( + model=model, llm_router=llm_router, team_id=team_id, team_model_aliases=team_model_aliases + ) + ) ## check model access for alias + underlying model - allow if either is in allowed models for m in potential_models: @@ -4434,6 +4459,11 @@ def _can_object_call_model( team_id=team_id, ): return True + if any( + _check_model_access_helper(model=name, llm_router=llm_router, models=models, team_id=team_id) + for name in owning_team_names + ): + return True internal_message: Final = ( f"{object_type} not allowed to access model. This {object_type} can only access models={models}. " diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index aadb5881ce0..03817bc7819 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -30,6 +30,7 @@ from litellm.proxy.auth.auth_utils import ( log_once_if_budget_reservation_disabled, warn_once_if_custom_auth_skips_common_checks, ) +from litellm.router import Router class TestCustomAuthCommonChecksWarning: @@ -1167,6 +1168,124 @@ async def test_managed_batch_routes_pass_team_model_access_check(route, request_ ) +_TEAM_DEPLOYMENT_ID: Final = "team-a-deployment-id" +_TEAM_INTERNAL_MODEL_NAME: Final = "model_name_team-a_6d1f3c2e-0000-4000-8000-000000000000" +_TEAM_PUBLIC_MODEL_NAME: Final = "team-batch-model" +_TEAM_MANAGED_BATCH_ID: Final = _encode_managed_id( + f"litellm_proxy;model_id:{_TEAM_DEPLOYMENT_ID};llm_batch_id:provider-batch-789" +) + + +def _team_scoped_router(public_model_name: str | None = _TEAM_PUBLIC_MODEL_NAME) -> Router: + return Router( + model_list=[ + { + "model_name": _TEAM_INTERNAL_MODEL_NAME, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + "model_info": { + "id": _TEAM_DEPLOYMENT_ID, + "team_id": "team-a", + **({"team_public_model_name": public_model_name} if public_model_name else {}), + }, + }, + ] + ) + + +def _team_scoped_character_id() -> str: + from litellm.types.videos.utils import encode_character_id_with_provider + + return encode_character_id_with_provider(character_id="char_team", provider="openai", model_id=_TEAM_DEPLOYMENT_ID) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route, request_data", + [ + pytest.param("/v1/batches/{batch_id}", {"batch_id": _TEAM_MANAGED_BATCH_ID}, id="batch-retrieve"), + pytest.param("/v1/batches/{batch_id}/cancel", {"batch_id": _TEAM_MANAGED_BATCH_ID}, id="batch-cancel"), + pytest.param( + "/v1/fine_tuning/jobs/{fine_tuning_job_id}", + { + "fine_tuning_job_id": _encode_managed_id( + f"litellm_proxy;model_id:{_TEAM_DEPLOYMENT_ID};generic_response_id:ftjob-1" + ) + }, + id="fine-tuning-job", + ), + pytest.param( + "/v1/vector_stores/{vector_store_id}", + { + "vector_store_id": _encode_managed_id( + "litellm_proxy:vector_store;unified_id,vs-unified;" + f"target_model_names,{_TEAM_PUBLIC_MODEL_NAME};resource_id,vs_1;model_id,{_TEAM_DEPLOYMENT_ID}" + ) + }, + id="vector-store", + ), + pytest.param( + "/v1/videos/characters/{character_id}", + {"character_id": _team_scoped_character_id()}, + id="character", + ), + ], +) +async def test_team_scoped_managed_ids_authorize_only_the_owning_team_by_public_name( + route: str, request_data: dict[str, str] +) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_checks import can_key_call_resolved_model + + llm_router: Final = _team_scoped_router() + + async def check(team_id: str) -> None: + await can_key_call_resolved_model( + model=get_model_from_request( + request_data=request_data, route=route, llm_router=llm_router, team_id=team_id + ), + llm_model_list=None, + valid_token=UserAPIKeyAuth(api_key="sk-team", team_id=team_id, team_models=[_TEAM_PUBLIC_MODEL_NAME]), + llm_router=llm_router, + ) + + await check("team-a") + with pytest.raises(ProxyException) as exc_info: + await check("team-b") + assert exc_info.value.code == "403" + + +@pytest.mark.asyncio +async def test_legacy_team_alias_authorizes_a_team_batch_only_for_keys_granted_the_alias() -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_checks import can_key_call_resolved_model + + llm_router: Final = _team_scoped_router(public_model_name=None) + + async def check(key_models: list[str]) -> None: + await can_key_call_resolved_model( + model=get_model_from_request( + request_data={"batch_id": _TEAM_MANAGED_BATCH_ID}, + route="/v1/batches/{batch_id}", + llm_router=llm_router, + team_id="team-a", + ), + llm_model_list=None, + valid_token=UserAPIKeyAuth( + api_key="sk-team", + team_id="team-a", + models=key_models, + team_models=["legacy-team-model"], + team_model_aliases={"legacy-team-model": _TEAM_INTERNAL_MODEL_NAME}, + ), + llm_router=llm_router, + ) + + await check(["legacy-team-model"]) + with pytest.raises(ProxyException) as exc_info: + await check(["some-other-model"]) + assert exc_info.value.code == "403" + + def test_get_model_from_request_resolves_character_id_model_with_router(): from litellm.types.videos.utils import encode_character_id_with_provider diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 2d597abf3b8..ee2056b0b98 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -3258,6 +3258,129 @@ async def test_cancel__unified_batch_id_rejects_key_without_model_grant(cancel_h cancel_harness.router_acancel.assert_not_called() +TEAM_INTERNAL_MODEL_NAME: Final = "model_name_team-a_00000000-0000-0000-0000-000000000000" +TEAM_PUBLIC_MODEL_NAME: Final = "team-gpt4o-mini" + + +def _team_scoped_deployment() -> Deployment: + return Deployment( + model_name=TEAM_INTERNAL_MODEL_NAME, + litellm_params=LiteLLM_Params(model="gpt-4o-mini"), + model_info=ModelInfo(id="dep-team-a-1", team_id="team-a", team_public_model_name=TEAM_PUBLIC_MODEL_NAME), + ) + + +def _route_unified_batch_to_team_deployment( + harness: RetrieveHarness | CancelHarness, deployment: Deployment | None = None +) -> None: + target: Final = deployment or _team_scoped_deployment() + harness.router.resolve_model_name_from_model_id = MagicMock(return_value=target.model_name) + harness.router.get_deployment_by_model_group_name = MagicMock(return_value=target) + + +@pytest.mark.asyncio +async def test_retrieve__team_scoped_deployment_authorizes_owning_team_by_public_name( + retrieve_harness: RetrieveHarness, +) -> None: + _route_unified_batch_to_team_deployment(retrieve_harness) + + resp: Final = await call_retrieve( + retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to(TEAM_PUBLIC_MODEL_NAME) + ) + + assert retrieve_harness.router_aretrieve.call_count == 1 + assert resp._hidden_params["model_id"] == "gpt-4o-mini" + + +@pytest.mark.asyncio +async def test_cancel__team_scoped_deployment_authorizes_owning_team_by_public_name( + cancel_harness: CancelHarness, +) -> None: + _route_unified_batch_to_team_deployment(cancel_harness) + + await call_cancel(cancel_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to(TEAM_PUBLIC_MODEL_NAME)) + + assert cancel_harness.router_acancel.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment, user", + [ + pytest.param( + _team_scoped_deployment(), _key_restricted_to("some-other-model"), id="owning-team-key-without-the-model" + ), + pytest.param( + _team_scoped_deployment(), + UserAPIKeyAuth( + api_key="sk-team-b", + team_id="team-b", + team_models=[TEAM_PUBLIC_MODEL_NAME], + models=[TEAM_PUBLIC_MODEL_NAME], + ), + id="other-team-key-granted-the-same-public-name", + ), + pytest.param( + Deployment( + model_name="model_name_team-a_ml_00000000-0000-0000-0000-000000000000", + litellm_params=LiteLLM_Params(model="gpt-4o-mini"), + model_info=ModelInfo( + id="dep-team-a-ml-1", team_id="team-a_ml", team_public_model_name=TEAM_PUBLIC_MODEL_NAME + ), + ), + _key_restricted_to(TEAM_PUBLIC_MODEL_NAME), + id="team-whose-id-prefixes-the-owning-team-id", + ), + ], +) +async def test_retrieve__team_scoped_deployment_rejects_keys_outside_its_grant( + retrieve_harness: RetrieveHarness, deployment: Deployment, user: UserAPIKeyAuth +) -> None: + _route_unified_batch_to_team_deployment(retrieve_harness, deployment) + + with pytest.raises(ProxyException) as exc_info: + await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=user) + + assert exc_info.value.code == "403" + retrieve_harness.router_aretrieve.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment, user", + [ + pytest.param( + _team_scoped_deployment(), + _key_restricted_to(TEAM_INTERNAL_MODEL_NAME), + id="owning-team-key-granted-the-internal-name", + ), + pytest.param( + Deployment( + model_name=TEAM_INTERNAL_MODEL_NAME, + litellm_params=LiteLLM_Params(model="gpt-4o-mini"), + model_info=ModelInfo(id="dep-team-a-1", team_id="team-a"), + ), + UserAPIKeyAuth( + api_key="sk-legacy", + team_id="team-a", + team_models=["legacy-team-model"], + models=["legacy-team-model"], + team_model_aliases={"legacy-team-model": TEAM_INTERNAL_MODEL_NAME}, + ), + id="legacy-team-model-reached-through-a-team-alias", + ), + ], +) +async def test_retrieve__team_scoped_deployment_authorizes_internal_name_and_legacy_alias_grants( + retrieve_harness: RetrieveHarness, deployment: Deployment, user: UserAPIKeyAuth +) -> None: + _route_unified_batch_to_team_deployment(retrieve_harness, deployment) + + await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=user) + + assert retrieve_harness.router_aretrieve.call_count == 1 + + @pytest.mark.asyncio async def test_cancel__executed_batch_rejects_key_without_model_grant(cancel_harness, executed_runner): runner, factory = executed_runner From d11f438e3defb4e40512ee360e41ee586069daa2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:58:56 -0700 Subject: [PATCH 02/36] bump: version 1.104.4 --- pyproject.toml | 4 ++-- uv.lock | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b3740e7f753..2816de2316c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.104.3" +version = "1.104.4" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -355,7 +355,7 @@ litellm-enterprise = { workspace = true } members = ["enterprise", "litellm-proxy-extras"] [tool.commitizen] -version = "1.104.3" +version = "1.104.4" version_files = [ "pyproject.toml:^version", ] diff --git a/uv.lock b/uv.lock index 527604aa4ed..fb9d8bbd9ac 100644 --- a/uv.lock +++ b/uv.lock @@ -4500,7 +4500,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.104.3" +version = "1.104.4" source = { editable = "." } dependencies = [ { name = "aiohttp" }, From c47b760253d84febfac19c63757b8d06e11b1416 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 09:10:10 -0700 Subject: [PATCH 03/36] test(proxy): import Final in the line's batch endpoint tests The tests picked from #47067 annotate module constants with Final, which main's copy imports and this line's did not, so the module failed to import on Python 3.12. --- tests/test_litellm/proxy/batches_endpoints/test_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index ee2056b0b98..27ce59af56e 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -35,7 +35,7 @@ import logging from contextlib import ExitStack from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx From 9d0cca7ed0d841938540cd7b4a1e7fda08819505 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:09:19 -0700 Subject: [PATCH 04/36] fix(token_counter): count Gemini function_declarations tools (#43417) Backport of #43417 to stable/1.104.x. Cherry-picked from 8e6d99d74a63c61e39628baeabea0c07dbeda5f4 (main). Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/token_counter.py | 78 +++++++++++-------- .../litellm_core_utils/test_token_counter.py | 72 +++++++++++++++++ 2 files changed, 119 insertions(+), 31 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 5d7956059e4..cdd2d0654be 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -951,7 +951,7 @@ def _count_content_list( ) -def _format_function_definitions(tools): +def _format_function_definitions(tools: Sequence[object]) -> str: """Formats tool definitions in the format that OpenAI appears to use. Based on https://github.com/forestwanglin/openai-java/blob/main/jtokkit/src/main/java/xyz/felh/openai/jtokkit/utils/TikTokenUtils.java """ @@ -959,41 +959,57 @@ def _format_function_definitions(tools): lines.append("namespace functions {") lines.append("") for tool in tools: - if not isinstance(tool, dict): + if not isinstance(tool, Mapping): continue - function = tool.get("function") - if not isinstance(function, dict): - # Anthropic tool shape → OpenAI function dict for token counting. - params = tool.get("input_schema") or tool.get("parameters") or {} - if not isinstance(params, dict): - params = {} - function = { - "name": tool.get("name"), - "description": tool.get("description"), - "parameters": params, - } - function_name = function.get("name") - if not function_name: - # Skip malformed tools missing a name to avoid emitting - # ``type None = ...`` which would produce inaccurate token counts. - continue - if function_description := function.get("description"): - lines.append(f"// {function_description}") - parameters = function.get("parameters") or {} - if not isinstance(parameters, dict): - parameters = {} - properties = parameters.get("properties") - if properties and properties.keys(): - lines.append(f"type {function_name} = (_: {{") - lines.append(_format_object_parameters(parameters, 0)) - lines.append("}) => any;") - else: - lines.append(f"type {function_name} = () => any;") - lines.append("") + for function in _function_definitions_for_tool(cast(Mapping[str, object], tool)): + lines.extend(_format_single_function_definition(function)) lines.append("} // namespace functions") return "\n".join(lines) +def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mapping[str, object]]: + function: Final = tool.get("function") + if isinstance(function, Mapping): + yield function + return + declarations: Final = tool.get("function_declarations") or tool.get("functionDeclarations") + if isinstance(declarations, list): + for declaration in declarations: + if isinstance(declaration, Mapping): + yield declaration + return + parameters: Final = tool.get("input_schema") or tool.get("parameters") or {} + normalized_parameters: Final = parameters if isinstance(parameters, Mapping) else {} + yield { + "name": tool.get("name"), + "description": tool.get("description"), + "parameters": normalized_parameters, + } + + +def _format_single_function_definition(function: Mapping[str, object]) -> tuple[str, ...]: + function_name: Final = function.get("name") + if not function_name: + return () + function_description: Final = function.get("description") + parameters_value: Final = function.get("parameters") or {} + parameters: Final = parameters_value if isinstance(parameters_value, Mapping) else {} + properties: Final = parameters.get("properties") + if isinstance(properties, Mapping) and properties: + return ( + *((f"// {function_description}",) if function_description else ()), + f"type {function_name} = (_: {{", + _format_object_parameters(parameters, 0), + "}) => any;", + "", + ) + return ( + *((f"// {function_description}",) if function_description else ()), + f"type {function_name} = () => any;", + "", + ) + + def _format_object_parameters(parameters, indent): properties: Final = parameters.get("properties") if not properties: diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index f7ded4f3fa8..c71b1496bdd 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -442,6 +442,78 @@ def test_token_counter_with_tools(message_count_pair): ), f"Expected {expected_tokens} tokens, got {counted_tokens}." +def test_token_counter_counts_gemini_function_declarations(): + openai_tools: Final = [ + { + "type": "function", + "function": { + "name": "lookup_weather", + "description": "Find current weather conditions for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City and region"}, + "units": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["location"], + }, + }, + } + ] + gemini_tools: Final = litellm.utils.get_optional_params( + model="gemini-2.5-pro", + custom_llm_provider="gemini", + tools=openai_tools, + )["tools"] + camel_case_tools: Final = [{"functionDeclarations": gemini_tools[0]["function_declarations"]}] + + openai_tokens: Final = token_counter_new( + model="gemini-2.5-pro", + messages=[{"role": "user", "content": "What's the weather?"}], + tools=openai_tools, + ) + gemini_tokens: Final = token_counter_new( + model="gemini-2.5-pro", + messages=[{"role": "user", "content": "What's the weather?"}], + tools=gemini_tools, + ) + camel_case_tokens: Final = token_counter_new( + model="gemini-2.5-pro", + messages=[{"role": "user", "content": "What's the weather?"}], + tools=camel_case_tools, + ) + + assert openai_tokens == gemini_tokens == camel_case_tokens + + +def test_token_counter_skips_non_mapping_tools(): + openai_tool: Final = { + "type": "function", + "function": { + "name": "lookup_weather", + "description": "Find current weather conditions for a location", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string", "description": "City and region"}}, + "required": ["location"], + }, + }, + } + messages: Final = [{"role": "user", "content": "What's the weather?"}] + valid_tokens: Final = token_counter_new( + model="gemini-2.5-pro", + messages=messages, + tools=[openai_tool], + ) + mixed_tokens: Final = token_counter_new( + model="gemini-2.5-pro", + messages=messages, + tools=["bad", None, openai_tool], + ) + + assert mixed_tokens == valid_tokens + + class NeedsToleranceUpdateError(Exception): """Custom exception to mark tests that have improved""" From a6a1956a1774aec0732350960a25c0a957a47291 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:09:19 -0700 Subject: [PATCH 05/36] fix(vertex_ai): forward system and tools to partner model count_tokens (#43900) Backport of #43900 to stable/1.104.x. Cherry-picked from 6d73fa6b491089f14a4ebb679abd0a835fa5182a (main). Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: jesus Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/vertex_ai/common_utils.py | 37 +- .../vertex_ai_partner_models/main.py | 6 + tests/integration/_support/vertex.py | 31 + .../test_vertex_partner_count_tokens_wire.py | 728 ++++++++++++++++++ .../test_vertex_partner_count_tokens_sdk.py | 90 +++ .../vertex_ai/test_vertex_ai_common_utils.py | 143 ++++ 6 files changed, 1027 insertions(+), 8 deletions(-) create mode 100644 tests/integration/_support/vertex.py create mode 100644 tests/integration/providers/test_vertex_partner_count_tokens_wire.py create mode 100644 tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 6d050d5a856..ead96da5cea 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1240,6 +1240,9 @@ class VertexAITokenCounter(BaseTokenCounter): ) -> TokenCountResponse | None: import copy + from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( + VertexAIError as PartnerVertexAIError, + ) from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( VertexAIPartnerModels, ) @@ -1269,14 +1272,32 @@ class VertexAITokenCounter(BaseTokenCounter): "vertex_ai_credentials" ) - result = await partner_models_handler.count_tokens( - model=model_to_use, - messages=messages or [], - litellm_params=partner_litellm_params, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_credentials=vertex_credentials, - ) + try: + result = await partner_models_handler.count_tokens( + model=model_to_use, + messages=messages or [], + litellm_params=partner_litellm_params, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_credentials=vertex_credentials, + system=system, + tools=tools, + ) + except (PartnerVertexAIError, httpx.HTTPStatusError) as e: + status_code: Final = e.response.status_code + error_message: Final = e.message if isinstance(e, PartnerVertexAIError) else e.response.text + verbose_logger.warning( + "Vertex AI partner CountTokens API error: status=%s, message=%s", status_code, error_message + ) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="vertex_ai_partner_models", + error=True, + error_message=error_message, + status_code=status_code, + ) if result is not None: return TokenCountResponse( diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index 2a36e5cc785..945f1808f7f 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -2,6 +2,7 @@ ## API Handler for calling Vertex AI Partner Models from collections.abc import Callable from enum import Enum +from types import MappingProxyType from typing import Final import httpx @@ -270,6 +271,8 @@ class VertexAIPartnerModels(VertexBase): vertex_project=None, vertex_location=None, vertex_credentials=None, + system: object | None = None, + tools: list[dict[str, object]] | None = None, ): """ Count tokens for Vertex AI partner models (Anthropic Claude, Mistral, etc.) @@ -303,6 +306,9 @@ class VertexAIPartnerModels(VertexBase): request_data: Final = { "model": model, "messages": messages, + **MappingProxyType( + {key: value for key, value in (("system", system), ("tools", tools)) if value is not None} + ), } # Prepare litellm_params with credentials diff --git a/tests/integration/_support/vertex.py b/tests/integration/_support/vertex.py new file mode 100644 index 00000000000..af7178fa6e9 --- /dev/null +++ b/tests/integration/_support/vertex.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import json +from typing import Final + +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa + + +def service_account_json(project: str, token_url: str) -> str: + private_key: Final = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": project, + "private_key_id": "scripted", + "private_key": private_key, + "client_email": f"scripted@{project}.iam.gserviceaccount.com", + "client_id": "0", + "auth_uri": f"{token_url}/_oauth/authorize", + "token_uri": f"{token_url}/_oauth/token", + } + ) diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py new file mode 100644 index 00000000000..5df8dd6e32e --- /dev/null +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -0,0 +1,728 @@ +import json +import re +import signal +import socket +import threading +import uuid +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from pathlib import Path +from queue import SimpleQueue +from typing import Final + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually +from integration._support.process import owned_proxy_process +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "claude-sonnet-4-6" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_MODELS_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models" +_COUNT_TARGET: Final = f"{_MODELS_PATH}/count-tokens:rawPredict" +_MESSAGE_TARGET: Final = f"{_MODELS_PATH}/{_BACKEND}:rawPredict" +_STREAM_TARGET: Final = f"{_MODELS_PATH}/{_BACKEND}:streamRawPredict?alt=sse" +_PEER_COUNT: Final = 4242 +_REJECTION: Final = "scripted partner rejection" +_REJECT_TEXT: Final = "The peer must reject this message" +_REPLY_TEXT: Final = "scripted reply" +_OWNED_MODEL: Final = "partner-claude" +_OWNED_UNREACHABLE_MODEL: Final = "partner-claude-unreachable" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") + +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this message"}] +_SYSTEM: Final = "You are a terse assistant that answers in one sentence" +_SYSTEM_BLOCKS: Final[list[JsonValue]] = [ + {"type": "text", "text": "You are a terse assistant"}, + {"type": "text", "text": "Answer in one sentence"}, +] +_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string", "description": "City to look up"}}, + "required": ["city"], +} +_TOOLS: Final[list[JsonValue]] = [ + {"name": "get_weather", "description": "Look up the current weather for a city", "input_schema": _WEATHER_SCHEMA} +] +_OPENAI_TOOLS: Final[list[JsonValue]] = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Look up the current weather for a city", + "parameters": _WEATHER_SCHEMA, + }, + } +] +_RESPONSES_TOOLS: Final[list[JsonValue]] = [ + { + "type": "function", + "name": "get_weather", + "description": "Look up the current weather for a city", + "parameters": _WEATHER_SCHEMA, + } +] +_PEER_BARE: Final[dict[str, JsonValue]] = {"model": _BACKEND, "messages": _MESSAGES} +_PEER_FULL: Final[dict[str, JsonValue]] = {**_PEER_BARE, "system": _SYSTEM, "tools": _TOOLS} +_GEMINI_BODY: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "Count this"}]}]} +_GEMINI_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this"}] + +_REPLY: Final[dict[str, JsonValue]] = { + "id": "msg_scripted", + "type": "message", + "role": "assistant", + "model": _BACKEND, + "content": [{"type": "text", "text": _REPLY_TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, +} +_EVENTS: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = ( + ( + "message_start", + { + "type": "message_start", + "message": {**_REPLY, "content": [], "stop_reason": None, "usage": {"input_tokens": 5, "output_tokens": 1}}, + }, + ), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _REPLY_TEXT}}, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + ), + ("message_stop", {"type": "message_stop"}), +) +_SSE: Final = tuple(f"event: {name}\ndata: {json.dumps(data)}\n\n".encode() for name, data in _EVENTS) + + +def _counted(_request: Request) -> Reply: + return Reply(body=json.dumps({"input_tokens": _PEER_COUNT}).encode()) + + +def _rejected(status: int) -> Reply: + return Reply( + status=status, + body=json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": _REJECTION}}).encode(), + ) + + +def _rejecting(status: int) -> Callable[[Request], Reply]: + def count(_request: Request) -> Reply: + return _rejected(status) + + return count + + +def _anthropic_message(message: JsonValue) -> bool: + return isinstance(message, dict) and message.get("role") in ("user", "assistant") + + +def _anthropic_tool(tool: JsonValue) -> bool: + return isinstance(tool, dict) and isinstance(tool.get("name"), str) and isinstance(tool.get("input_schema"), dict) + + +def _strict(request: Request) -> Reply: + body: Final = _JSON_OBJECT.validate_json(request.body) + messages: Final = body.get("messages") + tools: Final = body.get("tools", []) + accepted: Final = ( + isinstance(messages, list) + and all(map(_anthropic_message, messages)) + and isinstance(body.get("system", ""), (str, list)) + and isinstance(tools, list) + and all(map(_anthropic_tool, tools)) + ) + return _counted(request) if accepted else _rejected(400) + + +def _rejecting_marked_messages(request: Request) -> Reply: + return _rejected(400) if _REJECT_TEXT in request.body.decode() else _counted(request) + + +def _peer(count: Callable[[Request], Reply] = _counted) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target.endswith("/count-tokens:rawPredict"): + return count(request) + if request.target.endswith(":streamRawPredict?alt=sse"): + return Reply(content_type="text/event-stream", chunks=_SSE) + if request.target.endswith(f"/{_BACKEND}:rawPredict"): + return Reply(body=json.dumps(_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + return respond + + +def _count_requests(requests: Sequence[Request]) -> tuple[Request, ...]: + return tuple(request for request in requests if "count-tokens" in request.target) + + +def _count_bodies(requests: Sequence[Request], target: str = _COUNT_TARGET) -> tuple[dict[str, JsonValue], ...]: + counts: Final = _count_requests(requests) + for request in counts: + assert (request.method, request.target) == ("POST", target), request.target + assert request.headers["authorization"] == "Bearer scripted-token", request.headers + return tuple(_JSON_OBJECT.validate_json(request.body) for request in counts) + + +def _counted_bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]: + return _count_bodies(wire.drain()) + + +def _bare(model: str) -> dict[str, JsonValue]: + return {"model": model, "messages": _MESSAGES} + + +def _full(model: str) -> dict[str, JsonValue]: + return {**_bare(model), "system": _SYSTEM, "tools": _TOOLS} + + +def _deployment(gateway: Gateway, scenario: Scenario, api_base: str, **overrides: JsonValue) -> str: + return scenario.model( + **{ + "model": f"vertex_ai/{_BACKEND}", + "api_base": api_base, + "api_key": None, + "vertex_project": _PROJECT, + "vertex_location": _LOCATION, + "vertex_credentials": service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + **overrides, + } + ) + + +def _count(gateway: Gateway, body: Mapping[str, JsonValue]) -> httpx.Response: + return gateway.request("POST", "/v1/messages/count_tokens", body) + + +def _payload(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text + return _JSON_OBJECT.validate_json(response.content) + + +def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int: + response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "false"}) + payload: Final = _payload(response) + total: Final = payload["total_tokens"] + assert payload["tokenizer_type"] != "vertex_ai_partner_models", response.text + assert isinstance(total, int) and 0 < total != _PEER_COUNT, response.text + return total + + +def _closed_port_url() -> str: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + return f"http://127.0.0.1:{reserve.getsockname()[1]}" + + +def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]: + return tuple( + stack.enter_context(httpx.Client(base_url=base_url, timeout=30, trust_env=False)) for _ in range(count) + ) + + +def _counted_on(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue]: + response: Final = client.post( + "/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {key}"} + ) + return response.status_code, _JSON_OBJECT.validate_json(response.content).get("input_tokens") + + +def _generated_then_counted( + client: httpx.Client, key: str, model: str, body: Mapping[str, JsonValue] +) -> tuple[int, int, JsonValue]: + generated: Final = client.post( + "/v1/messages", + json={ + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"Generate before counting {uuid.uuid4().hex}"}], + }, + headers={"Authorization": f"Bearer {key}"}, + ) + return generated.status_code, *_counted_on(client, key, body) + + +def _local_port(client: httpx.Client) -> int: + with client.stream("GET", "/health/liveliness") as response: + port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1]) + response.read() + assert response.status_code == 200, response.text + return port + + +def _counted_or_dropped(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue] | None: + try: + return _counted_on(client, key, body) + except httpx.TransportError: + return None + + +def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]: + return frozenset( + connection.raddr.port + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.raddr and connection.laddr.port == proxy_port + ) + + +def _owned_config( + path: Path, gateway: Gateway, api_bases: Mapping[str, str], settings: Mapping[str, JsonValue] +) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path.write_text( + yaml.safe_dump( + { + **config, + "model_list": [ + { + "model_name": name, + "litellm_params": { + "model": f"vertex_ai/{_BACKEND}", + "api_base": api_base, + "vertex_project": _PROJECT, + "vertex_location": _LOCATION, + "vertex_credentials": service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + }, + } + for name, api_base in api_bases.items() + ], + "litellm_settings": {**config["litellm_settings"], **settings}, + } + ) + ) + return path + + +@pytest.mark.parametrize("system", [_SYSTEM, _SYSTEM_BLOCKS], ids=["string", "blocks"]) +def test_messages_count_tokens_forwards_system_and_tools(gateway: Gateway, system: JsonValue) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) + assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system, "tools": _TOOLS},) + + +def test_messages_count_tokens_without_system_or_tools_sends_bare_body(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, _bare(model)) + assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _counted_bodies(wire) == (_PEER_BARE,) + + +def test_utils_token_counter_call_endpoint_forwards_system_and_tools(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request( + "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} + ) + payload: Final = _payload(response) + assert (payload["total_tokens"], payload["tokenizer_type"]) == (_PEER_COUNT, "vertex_ai_partner_models") + assert (payload["request_model"], payload["model_used"]) == (model, _BACKEND), response.text + assert _counted_bodies(wire) == (_PEER_FULL,) + + +def test_utils_token_counter_local_mode_never_calls_the_peer(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + assert _local_count(gateway, _full(model)) > 0 + assert wire.drain() == () + + +def test_utils_token_counter_falls_back_locally_when_peer_rejects_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 _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: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request( + "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"} + ) + assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _PEER_COUNT}, response.text + assert _counted_bodies(wire) == (_PEER_BARE,) + + +def test_responses_input_tokens_falls_back_locally_when_peer_rejects(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( + "POST", + "/v1/responses/input_tokens", + {"model": model, "input": "Count this message", "instructions": "Be terse", "tools": _RESPONSES_TOOLS}, + ) + payload: Final = _payload(response) + (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: + with wire_server(_peer()) 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) + assert "totalTokens" in _payload(response), response.text + assert _counted_bodies(wire) == ({"model": _BACKEND, "messages": _GEMINI_MESSAGES},) + + +def test_gemini_count_tokens_route_falls_back_locally_when_peer_rejects(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 _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: + 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 _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 respond(request: Request) -> Reply: + if request.target == "/_oauth/token": + return Reply( + status=400, + body=json.dumps({"error": "invalid_grant", "error_description": "scripted refusal"}).encode(), + ) + return _peer()(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment( + gateway, scenario, wire.url, vertex_credentials=service_account_json(_PROJECT, wire.url) + ) + response: Final = _count(gateway, _full(model)) + payload: Final = _payload(response) + 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: + 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 + + +@pytest.mark.parametrize( + "tools", + [5, "", "x" * 5120, ["get_weather"]], + ids=["int", "empty_string", "5kb_string", "list_of_strings"], +) +def test_messages_count_tokens_rejects_malformed_tools_without_calling_the_peer( + gateway: Gateway, tools: JsonValue +) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + refused: Final = _count(gateway, {**_bare(model), "tools": tools}) + assert 400 <= refused.status_code < 600, refused.text + assert "input_tokens" not in refused.text, refused.text + assert wire.drain() == () + assert _generated_then_counted(gateway.client, gateway.key, model, _bare(model)) == (200, 200, _PEER_COUNT) + assert _counted_bodies(wire) == (_PEER_BARE,) + + +def test_messages_count_tokens_forwards_an_empty_tools_list(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {**_bare(model), "tools": []}) + assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": []},) + + +def test_messages_count_tokens_forwards_an_empty_system_string(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {**_bare(model), "system": ""}) + assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _counted_bodies(wire) == ({**_PEER_BARE, "system": ""},) + + +def test_messages_count_tokens_leaves_null_system_and_tools_out(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {**_bare(model), "system": None, "tools": None}) + assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _counted_bodies(wire) == (_PEER_BARE,) + + +def test_messages_count_tokens_falls_back_locally_when_peer_rejects_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 _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: + system: Final = "Answer in one sentence. " * 214 + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {**_bare(model), "system": system}) + assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system},) + + +def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value_each(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + fields: Final = f'"system": {json.dumps(_SYSTEM)}, "tools": {json.dumps(_TOOLS)}' + response: Final = gateway.client.post( + "/v1/messages/count_tokens", + content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}', + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + (sent,) = _count_requests(wire.drain()) + assert _count_bodies((sent,)) == (_PEER_FULL,) + assert (sent.body.count(b'"system"'), sent.body.count(b'"tools"')) == (1, 1), sent.body + + +def test_messages_count_tokens_unauthenticated_request_never_reaches_the_peer(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", _full(model), key="sk-not-a-key-this-proxy-issued" + ) + assert response.status_code == 401, response.text + assert "input_tokens" not in response.text, response.text + assert wire.drain() == () + + +@pytest.mark.parametrize("fields", [{}, {"messages": []}], ids=["missing", "empty"]) +def test_messages_count_tokens_without_messages_is_rejected(gateway: Gateway, fields: dict[str, JsonValue]) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _count(gateway, {"model": model, "system": _SYSTEM, "tools": _TOOLS, **fields}) + assert response.status_code == 400, response.text + assert "messages parameter is required" in response.text, response.text + assert wire.drain() == () + + +def test_count_tokens_location_override_targets_the_count_region(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment( + gateway, scenario, wire.url, vertex_location="global", vertex_count_tokens_location="europe-west1" + ) + response: Final = _count(gateway, _bare(model)) + assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + target: Final = _COUNT_TARGET.replace(f"/locations/{_LOCATION}/", "/locations/europe-west1/") + assert _count_bodies(wire.drain(), target) == (_PEER_BARE,) + + +def test_messages_count_tokens_repeated_request_reaches_the_peer_each_time(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2)) + assert answers == ({"input_tokens": _PEER_COUNT},) * 2 + assert _counted_bodies(wire) == (_PEER_BARE,) * 2 + + +@pytest.mark.timeout(240) # boots an owned two-worker proxy with litellm_settings.disable_token_counter +def test_disabled_token_counter_surfaces_provider_failures_instead_of_counting_locally( + gateway: Gateway, tmp_path: Path +) -> None: + with wire_server(_peer(_rejecting_marked_messages)) as wire: + config: Final = _owned_config( + tmp_path / "disabled-token-counter.yaml", + gateway, + {_OWNED_MODEL: wire.url, _OWNED_UNREACHABLE_MODEL: _closed_port_url()}, + {"disable_token_counter": True}, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + counted: Final = _count(owned.gateway, _full(_OWNED_MODEL)) + assert _payload(counted) == {"input_tokens": _PEER_COUNT}, counted.text + rejected: Final = _count( + owned.gateway, {**_full(_OWNED_MODEL), "messages": [{"role": "user", "content": _REJECT_TEXT}]} + ) + assert rejected.status_code == 400, rejected.text + assert _REJECTION in rejected.text and "input_tokens" not in rejected.text, rejected.text + unreachable: Final = _count(owned.gateway, _full(_OWNED_UNREACHABLE_MODEL)) + assert 500 <= unreachable.status_code < 600, unreachable.text + assert "input_tokens" not in unreachable.text, unreachable.text + local: Final = owned.gateway.request( + "POST", "/utils/token_counter", _full(_OWNED_MODEL), params={"call_endpoint": "false"} + ) + assert local.status_code == 503, local.text + assert len(_counted_bodies(wire)) == 2 + + +def test_peer_outage_between_concurrent_waves_falls_back_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))) + scenario: Final = stack.enter_context(gateway.scenario()) + with wire_server(_peer()) as wire: + 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) + + 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) + 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) + + +def test_slow_peer_holds_concurrent_counts_without_stalling_the_proxy(gateway: Gateway) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + + def hold(request: Request) -> Reply: + held.put(request.target) + assert release.wait(timeout=20), "Held count was never released" + return _counted(request) + + with ExitStack() as stack: + clients: Final = _clients(stack, str(gateway.client.base_url), 6) + wire: Final = stack.enter_context(wire_server(_peer(hold))) + scenario: Final = stack.enter_context(gateway.scenario()) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + model: Final = _deployment(gateway, scenario, wire.url) + try: + futures: Final = tuple( + pool.submit(_generated_then_counted, client, gateway.key, model, _full(model)) for client in clients + ) + eventually(held.qsize, lambda size: size == len(clients), seconds=30) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + assert _local_count(gateway, _full(model)) > 0 + assert not any(future.done() for future in futures) + finally: + release.set() + assert tuple(future.result(timeout=30) for future in futures) == ((200, 200, _PEER_COUNT),) * len(clients) + assert len(_counted_bodies(wire)) == len(clients) + + +@pytest.mark.timeout(300) # boots an owned two-worker proxy, kills one worker, and waits for its replacement +def test_worker_sigkill_mid_burst_leaves_the_sibling_counting(gateway: Gateway, tmp_path: Path) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + + def hold(request: Request) -> Reply: + held.put(request.target) + assert release.wait(timeout=60), "Held count was never released" + return _counted(request) + + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_peer(hold))) + stack.callback(release.set) + config: Final = _owned_config(tmp_path / "worker-kill.yaml", gateway, {_OWNED_MODEL: wire.url}, {}) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2)) + proxy_url: Final = owned.gateway.client.base_url + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + clients: Final = _clients(stack, str(proxy_url), 12) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + stack.callback(release.set) + ports: Final = tuple(_local_port(client) for client in clients) + futures: Final = tuple( + pool.submit(_counted_or_dropped, client, gateway.key, _full(_OWNED_MODEL)) for client in clients + ) + eventually(held.qsize, lambda size: size == len(clients), seconds=30) + shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers} + assert sum(map(len, shares.values())) == len(clients), shares + victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) + psutil.Process(victim).send_signal(signal.SIGKILL) + release.set() + results: Final = tuple(future.result(timeout=60) for future in futures) + for port, result in zip(ports, results, strict=True): + assert result == (None if port in shares[victim] else (200, _PEER_COUNT)), (port, result, shares) + second_wave: Final = _clients(stack, str(proxy_url), 6) + assert tuple(_counted_on(client, gateway.key, _full(_OWNED_MODEL)) for client in second_wave) == ( + (200, _PEER_COUNT), + ) * len(second_wave) + assert len(_counted_bodies(wire)) == len(clients) + len(second_wave) + eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120) + assert owned.process.poll() is None + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_chat_completions_on_the_same_deployment_still_generate(gateway: Gateway, stream: bool) -> None: + marker: Final = f"chat control {uuid.uuid4().hex}" + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + assert not stream or response.text.rstrip().endswith("data: [DONE]"), response.text + (sent,) = wire.drain() + assert sent.target == (_STREAM_TARGET if stream else _MESSAGE_TARGET), sent.target + assert marker in sent.body.decode(), sent.body + + +def test_messages_endpoint_on_the_same_deployment_still_generates(gateway: Gateway) -> None: + marker: Final = f"messages control {uuid.uuid4().hex}" + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + ) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + (sent,) = wire.drain() + assert sent.target == _MESSAGE_TARGET, sent.target + assert marker in sent.body.decode(), sent.body + + +def test_responses_endpoint_on_the_same_deployment_still_generates(gateway: Gateway) -> None: + marker: Final = f"responses control {uuid.uuid4().hex}" + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": marker}) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + (sent,) = wire.drain() + assert sent.target == _MESSAGE_TARGET, sent.target + assert marker in sent.body.decode(), sent.body diff --git a/tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py b/tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py new file mode 100644 index 00000000000..e59d3f9264a --- /dev/null +++ b/tests/integration/sdk/test_vertex_partner_count_tokens_sdk.py @@ -0,0 +1,90 @@ +import json +from collections.abc import Callable +from typing import Final + +import litellm +import pytest +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "claude-sonnet-4-6" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_COUNT_TARGET: Final = ( + f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/count-tokens:rawPredict" +) +_PEER_COUNT: Final = 4242 +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGES: Final[list[dict[str, str]]] = [{"role": "user", "content": "Count this message"}] +_SYSTEM: Final = "You are a terse assistant that answers in one sentence" +_TOOLS: Final[list[dict[str, JsonValue]]] = [ + { + "name": "get_weather", + "description": "Look up the current weather for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } +] + + +def _peer(status: int) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == "/_oauth/token": + token: Final = {"access_token": "scripted-token", "token_type": "Bearer", "expires_in": 3600} + return Reply(body=json.dumps(token).encode()) + if status == 200: + return Reply(body=json.dumps({"input_tokens": _PEER_COUNT}).encode()) + rejection: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted"}} + return Reply(status=status, body=json.dumps(rejection).encode()) + + return respond + + +def _count_requests(requests: tuple[Request, ...]) -> tuple[Request, ...]: + return tuple(request for request in requests if "count-tokens" in request.target) + + +@pytest.fixture +def vertex_environment(monkeypatch: pytest.MonkeyPatch) -> Callable[[str], None]: + def configure(token_url: str) -> None: + monkeypatch.setenv("VERTEXAI_PROJECT", _PROJECT) + monkeypatch.setenv("VERTEXAI_LOCATION", _LOCATION) + monkeypatch.setenv("VERTEXAI_CREDENTIALS", service_account_json(_PROJECT, token_url)) + + return configure + + +async def test_acount_tokens_forwards_system_and_tools_to_the_partner_peer( + vertex_environment: Callable[[str], None], +) -> None: + with wire_server(_peer(200)) as wire: + vertex_environment(wire.url) + counted: Final = await litellm.acount_tokens( + model=f"vertex_ai/{_BACKEND}", messages=_MESSAGES, tools=_TOOLS, system=_SYSTEM, api_base=wire.url + ) + (sent,) = _count_requests(wire.drain()) + assert (sent.method, sent.target, sent.headers["authorization"]) == ( + "POST", + _COUNT_TARGET, + "Bearer scripted-token", + ) + assert _JSON_OBJECT.validate_json(sent.body) == { + "model": _BACKEND, + "messages": _MESSAGES, + "system": _SYSTEM, + "tools": _TOOLS, + } + assert (counted.total_tokens, counted.tokenizer_type) == (_PEER_COUNT, "vertex_ai_partner_models"), counted + + +async def test_acount_tokens_falls_back_to_the_local_tokenizer_when_the_peer_rejects( + vertex_environment: Callable[[str], None], +) -> None: + with wire_server(_peer(400)) as wire: + vertex_environment(wire.url) + counted: Final = await litellm.acount_tokens( + model=f"vertex_ai/{_BACKEND}", messages=_MESSAGES, tools=_TOOLS, system=_SYSTEM, api_base=wire.url + ) + assert len(_count_requests(wire.drain())) == 1 + assert counted.tokenizer_type == "local_tokenizer", counted + assert counted.total_tokens > 0 and counted.total_tokens != _PEER_COUNT, counted diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 04a7ee451c4..86eb26a4c15 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1229,6 +1229,149 @@ async def test_vertex_ai_token_counter_routes_partner_models(): assert result.tokenizer_type == "vertex_ai_partner_models" +@pytest.mark.asyncio +async def test_vertex_ai_token_counter_forwards_system_and_tools_to_partner_request(): + from typing import Final + from unittest.mock import AsyncMock, patch + + from litellm.llms.vertex_ai.common_utils import VertexAITokenCounter + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler + + class FakeResponse: + status_code = 200 + + def json(self) -> dict[str, int]: + return {"input_tokens": 37} + + class FakeHttpClient: + posted_bodies: tuple[dict[str, object], ...] = () + + async def post( + self, + url: str, + headers: dict[str, str], + json: dict[str, object], + timeout: float, + ) -> FakeResponse: + self.posted_bodies = (*self.posted_bodies, json) + return FakeResponse() + + fake_http_client: Final = FakeHttpClient() + counter: Final = VertexAITokenCounter() + model: Final = "claude-opus-5-5" + messages: Final = [{"role": "user", "content": "Hello"}] + system: Final = "Follow the system instructions" + tools: Final = [ + { + "name": "lookup", + "description": "Look up a value", + "input_schema": {"type": "object", "properties": {}}, + } + ] + deployment: Final = { + "litellm_params": { + "vertex_project": "test-project", + "vertex_location": "us-east5", + } + } + + with ( + patch.object(handler, "get_async_httpx_client", return_value=fake_http_client), + patch.object( + handler.VertexAIPartnerModelsTokenCounter, + "_ensure_access_token_async", + new=AsyncMock(return_value=("fake-token", "test-project")), + ), + ): + with_optional_fields: Final = await counter.count_tokens( + model_to_use=model, + messages=messages, + contents=None, + deployment=deployment, + system=system, + tools=tools, + ) + without_optional_fields: Final = await counter.count_tokens( + model_to_use=model, + messages=messages, + contents=None, + deployment=deployment, + ) + + assert fake_http_client.posted_bodies == ( + {"model": model, "messages": messages, "system": system, "tools": tools}, + {"model": model, "messages": messages}, + ) + assert with_optional_fields is not None + assert with_optional_fields.total_tokens == 37 + assert without_optional_fields is not None + assert without_optional_fields.total_tokens == 37 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("provider_failure", "expected_status", "expected_message"), + [ + ("http_400", 400, 'tools.0: Input tag "function" does not match any of the expected tags'), + ("credentials", 500, "could not resolve credentials"), + ], +) +async def test_vertex_ai_token_counter_returns_partner_provider_error_as_value( + provider_failure: str, expected_status: int, expected_message: str +): + from typing import Final + from unittest.mock import AsyncMock, patch + + import httpx + + from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError + from litellm.llms.vertex_ai.common_utils import VertexAITokenCounter + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler + + class RejectingHttpClient: + async def post( + self, + url: str, + headers: dict[str, str], + json: dict[str, object], + timeout: float, + ) -> None: + request: Final = httpx.Request("POST", url) + response: Final = httpx.Response(400, text=expected_message, request=request) + raise MaskedHTTPStatusError( + httpx.HTTPStatusError("Client error '400 Bad Request'", request=request, response=response), + message=expected_message, + text=expected_message, + ) + + access_token: Final = ( + AsyncMock(side_effect=ValueError(expected_message)) + if provider_failure == "credentials" + else AsyncMock(return_value=("fake-token", "test-project")) + ) + with ( + patch.object(handler, "get_async_httpx_client", return_value=RejectingHttpClient()), + patch.object(handler.VertexAIPartnerModelsTokenCounter, "_ensure_access_token_async", new=access_token), + ): + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="claude-opus-5-5", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}}, + request_model="vertex-claude", + tools=[{"type": "function", "function": {"name": "lookup", "parameters": {}}}], + ) + + assert result is not None + assert result.error is True + assert result.status_code == expected_status + assert result.error_message is not None + assert expected_message in result.error_message + assert result.total_tokens == 0 + assert result.request_model == "vertex-claude" + assert result.tokenizer_type == "vertex_ai_partner_models" + + @pytest.mark.asyncio async def test_vertex_ai_token_counter_uses_count_tokens_location(): """ From 2f968f566b00f337608d436ec93e10ceaa0ead90 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 19 Oct 2026 10:12:14 -0700 Subject: [PATCH 06/36] fix(anthropic): normalize images for provider token counting (#45185) Backport of #45185 to stable/1.104.x. Cherry-picked from 6befb9ad7dba1835f4952cc1430f750a9f6fa12a (main). --- .../llms/anthropic/count_tokens/handler.py | 3 +- .../anthropic/count_tokens/transformation.py | 32 +++- .../anthropic/count_tokens/handler.py | 3 +- ...t_anthropic_count_tokens_transformation.py | 179 ++++++++++++++++++ 4 files changed, 214 insertions(+), 3 deletions(-) diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index dd2135f4918..0cecf62368c 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -12,6 +12,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.llms.anthropic.common_utils import AnthropicError from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, @@ -62,7 +63,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): verbose_logger.debug("Processing Anthropic CountTokens request for model: %s", model) # Transform request to Anthropic format - request_body: Final = self.transform_request_to_count_tokens( + request_body: Final = await asyncify(self.transform_request_to_count_tokens)( model=model, messages=messages, tools=tools, diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index fb12747cec0..8147496b402 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -11,11 +11,41 @@ from typing import Final from pydantic import JsonValue, TypeAdapter from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION +from litellm.types.llms.openai import ChatCompletionImageObject _COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue]) +_IMAGE_BLOCK: Final = TypeAdapter(ChatCompletionImageObject) COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config") +def _count_image(block: JsonValue) -> JsonValue: + if not isinstance(block, dict) or block.get("type") != "image_url": + return block + from litellm.litellm_core_utils.prompt_templates.factory import convert_to_anthropic_image_obj + + image_block: Final = _IMAGE_BLOCK.validate_python(block) + image_url: Final = image_block["image_url"] + source: Final = convert_to_anthropic_image_obj( + openai_image_url=image_url if isinstance(image_url, str) else image_url["url"], + format=image_url.get("format") if isinstance(image_url, dict) else None, + ) + image: Final = _COUNT_REQUEST.validate_python({"type": "image", "source": source}) + return {**{key: value for key, value in block.items() if key not in {"type", "image_url"}}, **image} + + +def _count_block(block: JsonValue) -> JsonValue: + if not isinstance(block, dict) or block.get("type") != "tool_result": + return _count_image(block) + content: Final = block.get("content") + if not isinstance(content, list): + return block + return {**block, "content": [_count_image(part) for part in content]} + + +def _count_content(content: JsonValue) -> JsonValue: + return [_count_block(block) for block in content] if isinstance(content, list) else content + + class AnthropicCountTokensConfig: """ Configuration and transformation logic for Anthropic CountTokens API. @@ -53,7 +83,7 @@ class AnthropicCountTokensConfig: MappingProxyType( { "model": model, - "messages": messages, + "messages": [{**message, "content": _count_content(message["content"])} for message in messages], **MappingProxyType( {key: value for key, value in (("system", system), ("tools", tools)) if value 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 36d5a56db0d..5cbde9ef52b 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -10,6 +10,7 @@ import httpx import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.llms.anthropic.common_utils import AnthropicError from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( AzureAIAnthropicCountTokensConfig, @@ -59,7 +60,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): verbose_logger.debug("Processing Azure AI Anthropic CountTokens request for model: %s", model) # Transform request to Anthropic format - request_body: Final = self.transform_request_to_count_tokens( + request_body: Final = await asyncify(self.transform_request_to_count_tokens)( model=model, messages=messages, tools=tools, diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index ddac561f337..18d57e71d28 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -1,7 +1,113 @@ +import asyncio +from base64 import b64encode +from copy import deepcopy +from threading import get_ident +from typing import Final +import httpx +import pytest +import respx +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, ) +from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( + AzureAIAnthropicCountTokensHandler, +) +from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( + AzureAIAnthropicCountTokensConfig, +) + + +@pytest.mark.parametrize( + "config_type", (AnthropicCountTokensConfig, AzureAIAnthropicCountTokensConfig) +) +@pytest.mark.parametrize( + ("image_url", "source"), + ( + ("data:image/png;base64,aW1hZ2U=", {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}), + ({"url": "data:image/png;base64,aW1hZ2U="}, {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}), + ( + {"url": "data:image/png;base64,aW1hZ2U=", "format": "image/jpeg", "detail": "high"}, + {"type": "base64", "media_type": "image/jpeg", "data": "aW1hZ2U="}, + ), + ), +) +def test_count_translates_openai_images_without_mutating_input( + config_type: type[AnthropicCountTokensConfig], + image_url: str | dict[str, JsonValue], + source: dict[str, JsonValue], +) -> None: + cache_control: Final[dict[str, JsonValue]] = {"type": "ephemeral"} + messages: Final[list[dict[str, JsonValue]]] = [{ + "role": "user", "content": [ + {"type": "text", "text": "Count this image"}, + {"type": "image_url", "image_url": image_url, "cache_control": cache_control}, + ], + }] + original: Final = deepcopy(messages) + result: Final = config_type().transform_request_to_count_tokens( + model="claude-opus-5-5", messages=messages + ) + + assert result == { + "model": "claude-opus-5-5", "messages": [{ + "role": "user", "content": [ + {"type": "text", "text": "Count this image"}, + {"type": "image", "source": source, "cache_control": cache_control}, + ], + }], + } + assert messages == original + + +@pytest.mark.parametrize( + "config_type", (AnthropicCountTokensConfig, AzureAIAnthropicCountTokensConfig) +) +def test_count_normalizes_nested_tool_images_and_preserves_native_fields( + config_type: type[AnthropicCountTokensConfig], +) -> None: + openai_image: Final[dict[str, JsonValue]] = { + "type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="} + } + native_image: Final[dict[str, JsonValue]] = { + "type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="} + } + assistant: Final[dict[str, JsonValue]] = {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "inspect screenshot", "signature": "fixture-signature"}, + {"type": "tool_use", "id": "read-1", "name": "Read", "input": {"content": [openai_image]}}, + ]} + tool_result: Final[dict[str, JsonValue]] = { + "type": "tool_result", "tool_use_id": "read-1", "is_error": False, + "content": [{"type": "text", "text": "Screenshot"}, native_image, openai_image], + "cache_control": {"type": "ephemeral"}, + } + text_result: Final[dict[str, JsonValue]] = {"type": "tool_result", "tool_use_id": "read-2", "content": "done"} + messages: Final[list[dict[str, JsonValue]]] = [ + assistant, {"role": "user", "content": [native_image, tool_result, text_result]} + ] + tools: Final[list[dict[str, JsonValue]]] = [{ + "name": "Read", "input_schema": {"type": "object", "examples": [openai_image]} + }] + system: Final[JsonValue] = [{"type": "text", "text": "policy", "cache_control": {"type": "ephemeral"}}] + options: Final[dict[str, JsonValue]] = { + "thinking": {"type": "adaptive"}, "tool_choice": {"type": "auto"}, "output_config": {"effort": "high"} + } + original: Final = deepcopy((messages, tools, system, options)) + result: Final = config_type().transform_request_to_count_tokens( + model="claude-opus-5-5", messages=messages, tools=tools, system=system, optional_params=options + ) + + assert result == { + "model": "claude-opus-5-5", "system": system, "tools": tools, **options, + "messages": [assistant, {"role": "user", "content": [native_image, { + **tool_result, "content": [{"type": "text", "text": "Screenshot"}, native_image, native_image] + }, text_result]}], + } + assert (messages, tools, system, options) == original def test_transform_basic_request(): @@ -88,3 +194,76 @@ def test_transform_no_system_no_tools(): assert "system" not in result assert "tools" not in result + + +@pytest.fixture +def httpx_transport_clients(monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + client_cache = getattr(litellm, "in_memory_llm_clients_cache", None) + if client_cache is not None: + client_cache.flush_cache() + yield + if client_cache is not None: + client_cache.flush_cache() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@pytest.mark.parametrize( + "handler_type", (AnthropicCountTokensHandler, AzureAIAnthropicCountTokensHandler) +) +@pytest.mark.parametrize("scheme", ("http", "https")) +@pytest.mark.parametrize("dict_url", (False, True)) +async def test_remote_image_fetch_keeps_counting_handler_event_loop_responsive( + handler_type: type[AnthropicCountTokensHandler] | type[AzureAIAnthropicCountTokensHandler], + scheme: str, + dict_url: bool, +) -> None: + loop: Final = asyncio.get_running_loop() + loop_thread: Final = get_ident() + witness: Final = asyncio.Event() + image_bytes: Final = b"\x89PNG\r\n\x1a\ncount-image" + image_url: Final = f"{scheme}://1.1.1.1/{handler_type.__name__}-{dict_url}.png" + model: Final = "claude-opus-5-5" + api_base: Final = "https://gateway.example/anthropic" + image: Final[dict[str, JsonValue]] = { + "type": "image_url", "image_url": {"url": image_url} if dict_url else image_url, + "cache_control": {"type": "ephemeral"}, + } + messages: Final[list[dict[str, JsonValue]]] = [{ + "role": "user", "content": [image, {"type": "tool_result", "tool_use_id": "read-1", "content": [image]}] + }] + original: Final = deepcopy(messages) + + async def run_witness() -> None: + witness.set() + + def image_response(_request: httpx.Request) -> httpx.Response: + assert get_ident() != loop_thread, "image fetch blocked the counting handler's event loop" + asyncio.run_coroutine_threadsafe(run_witness(), loop).result(timeout=5) + return httpx.Response(200, content=image_bytes, headers={"Content-Type": "image/png"}) + + with respx.mock: + image_route: Final = respx.get(image_url).mock(side_effect=image_response) + count_route: Final = respx.post(url__startswith=api_base).mock( + return_value=httpx.Response(200, json={"input_tokens": 7}) + ) + handler: Final = handler_type() + result: Final = await handler.handle_count_tokens_request( + model=model, messages=messages, api_base=api_base, api_key="test-key" + ) + + assert witness.is_set() + assert image_route.call_count == count_route.call_count == 1 + assert result == {"input_tokens": 7} + native_image: Final[dict[str, JsonValue]] = { + "type": "image", "source": { + "type": "base64", "media_type": "image/png", "data": b64encode(image_bytes).decode() + }, "cache_control": {"type": "ephemeral"}, + } + assert TypeAdapter(dict[str, JsonValue]).validate_json(count_route.calls.last.request.content) == { + "model": model, "messages": [{"role": "user", "content": [ + native_image, {"type": "tool_result", "tool_use_id": "read-1", "content": [native_image]} + ]}], + } + assert messages == original From 3da404a7c12fceee1cade251643dcfff330091c1 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:14:05 -0700 Subject: [PATCH 07/36] fix(bedrock_mantle): add is_mantle_claude_model for the count_tokens picks Hand-port of the is_mantle_claude_model helper from #43646 to stable/1.104.x, which #45317 imports. Taken from 01b4ffe16b (main). --- litellm/llms/bedrock_mantle/common_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index 9892ef6224e..ca637af0110 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -158,3 +158,7 @@ def mantle_base_segment(model: str | None, model_cost: dict) -> str: """ entry: Final = model_cost.get(f"bedrock_mantle/{split_mantle_region_prefix(model)[1]}", {}) if model else {} return "openai/v1" if entry.get("use_openai_responses_path") is True else "v1" + + +def is_mantle_claude_model(model: str) -> bool: + return "claude" in model.lower() From 8222299f1d99a20b2fe2a48bb9302af7fee5e719 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:12:15 -0700 Subject: [PATCH 08/36] fix(bedrock): count tokens on bedrock-mantle when bedrock-runtime cannot count a Claude model (#45317) Backport of #45317 to stable/1.104.x. Cherry-picked from 577d74c1ceb73c31c05f4677c6df3a68d53b924c (main). --- .../count_tokens/bedrock_token_counter.py | 182 +-- litellm/llms/bedrock/count_tokens/handler.py | 2 +- .../bedrock/count_tokens/mantle_handler.py | 99 ++ litellm/proxy/google_endpoints/endpoints.py | 24 +- .../test_bedrock_mantle_count_tokens_wire.py | 1072 +++++++++++++++++ .../test_google_api_endpoints.py | 28 + ...est_bedrock_mantle_count_tokens_handler.py | 109 ++ .../test_bedrock_token_counter.py | 140 +++ tests/unit/proxy/test_proxy_token_counter.py | 92 +- 9 files changed, 1612 insertions(+), 136 deletions(-) create mode 100644 litellm/llms/bedrock/count_tokens/mantle_handler.py create mode 100644 tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py create mode 100644 tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py create mode 100644 tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 2d02b152c61..ed61707f99e 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -3,27 +3,100 @@ Bedrock Token Counter implementation using the CountTokens API. """ from collections.abc import Mapping, Sequence +from dataclasses import dataclass from typing import Any, Final +from pydantic import JsonValue +from typing_extensions import assert_never + from litellm._logging import verbose_logger from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler +from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountResponse + +RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" +MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" + + +class _CountTokensReply(LiteLLMPydanticObjectBase): + input_tokens: int + + +@dataclass(frozen=True, slots=True) +class CountedTokens: + input_tokens: int + original_response: Mapping[str, JsonValue] + tokenizer_type: str + + +@dataclass(frozen=True, slots=True) +class CountTokensFailure: + status_code: int + message: str + tokenizer_type: str + + +CountTokensOutcome = CountedTokens | CountTokensFailure + + +def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool: + return ( + isinstance(outcome, CountTokensFailure) + and outcome.status_code == 400 + and is_mantle_claude_model(resolved_model) + ) class BedrockTokenCounter(BaseTokenCounter): - """Token counter implementation for AWS Bedrock provider using the CountTokens API.""" + """Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens + on bedrock-mantle for the Claude models bedrock-runtime answers 400 for""" + + def __init__( + self, + runtime_handler: BedrockCountTokensHandler | None = None, + mantle_handler: BedrockMantleCountTokensHandler | None = None, + client: AsyncHTTPHandler | None = None, + ) -> None: + self._runtime_handler: Final = runtime_handler or BedrockCountTokensHandler() + self._mantle_handler: Final = mantle_handler or BedrockMantleCountTokensHandler() + self._client: Final = client def should_use_token_counting_api( self, custom_llm_provider: str | None = None, ) -> bool: - """ - Returns True if we should use the Bedrock CountTokens API for token counting. - """ return custom_llm_provider == LlmProviders.BEDROCK.value + async def _count_with( + self, + handler: BedrockCountTokensHandler | BedrockMantleCountTokensHandler, + tokenizer_type: str, + request_data: dict[str, object], + litellm_params: dict[str, object], + resolved_model: str, + ) -> CountTokensOutcome: + try: + result: Final = await handler.handle_count_tokens_request( + request_data=request_data, + litellm_params=litellm_params, + resolved_model=resolved_model, + client=self._client, + ) + reply: Final = _CountTokensReply.model_validate(result) + except BedrockError as e: + verbose_logger.debug( + "%s CountTokens API error: status=%s, message=%s", tokenizer_type, e.status_code, e.message + ) + return CountTokensFailure(status_code=e.status_code, message=e.message, tokenizer_type=tokenizer_type) + except Exception as e: + verbose_logger.debug("Error calling %s CountTokens API: %s", tokenizer_type, e) + return CountTokensFailure(status_code=500, message=str(e), tokenizer_type=tokenizer_type) + return CountedTokens(input_tokens=reply.input_tokens, original_response=result, tokenizer_type=tokenizer_type) + async def count_tokens( self, model_to_use: str, @@ -34,81 +107,52 @@ class BedrockTokenCounter(BaseTokenCounter): tools: Sequence[Mapping[str, object]] | None = None, system: object | None = None, ) -> TokenCountResponse | None: - """ - Count tokens using AWS Bedrock's CountTokens API. - - This method calls the existing BedrockCountTokensHandler to make an API call - to Bedrock's token counting endpoint, bypassing the local tiktoken-based counting. - - Args: - model_to_use: The model identifier - messages: The messages to count tokens for - contents: Alternative content format (not used for Bedrock) - deployment: Deployment configuration containing litellm_params - request_model: The original request model name - - Returns: - TokenCountResponse with token count, or None if counting fails - """ if not messages: return None - deployment = deployment or {} - litellm_params: Final = deployment.get("litellm_params", {}) - - # Build request data in the format expected by BedrockCountTokensHandler + litellm_params: Final = (deployment or {}).get("litellm_params", {}) request_data: Final[dict[str, object]] = { "model": model_to_use, "messages": messages, + **({"tools": tools} if tools else {}), + **({"system": system} if system else {}), } - - if tools: - request_data["tools"] = tools - - if system: - request_data["system"] = system - - # Get the resolved model (strip prefixes like bedrock/, converse/, etc.) resolved_model: Final = get_bedrock_base_model(model_to_use) - try: - handler: Final = BedrockCountTokensHandler() - result: Final = await handler.handle_count_tokens_request( - request_data=request_data, - litellm_params=litellm_params, - resolved_model=resolved_model, + runtime: Final = await self._count_with( + self._runtime_handler, RUNTIME_TOKENIZER_TYPE, request_data, litellm_params, resolved_model + ) + outcome: Final = ( + await self._count_with( + self._mantle_handler, MANTLE_TOKENIZER_TYPE, request_data, litellm_params, resolved_model ) - - # Transform response to TokenCountResponse - if result is not None: + if _runtime_rejected_claude_model(runtime, resolved_model) + else runtime + ) + match outcome: + case CountedTokens(): return TokenCountResponse( - total_tokens=result.get("input_tokens", 0), + total_tokens=outcome.input_tokens, request_model=request_model, model_used=model_to_use, - tokenizer_type="bedrock_api", - original_response=result, + tokenizer_type=outcome.tokenizer_type, + original_response=dict(outcome.original_response), ) - except BedrockError as e: - verbose_logger.warning("Bedrock CountTokens API error: status=%s, message=%s", e.status_code, e.message) - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="bedrock_api", - error=True, - error_message=e.message, - status_code=e.status_code, - ) - except Exception as e: - verbose_logger.warning("Error calling Bedrock CountTokens API: %s", e) - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="bedrock_api", - error=True, - error_message=str(e), - status_code=500, - ) - - return None + case CountTokensFailure(): + verbose_logger.warning( + "%s CountTokens API error: status=%s, message=%s", + outcome.tokenizer_type, + outcome.status_code, + outcome.message, + ) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type=outcome.tokenizer_type, + error=True, + error_message=outcome.message, + status_code=outcome.status_code, + ) + case _: + assert_never(outcome) diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index d7fb510f057..be8243e48ef 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -125,7 +125,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): raise except httpx.HTTPStatusError as e: # HTTP errors - preserve the actual status code - verbose_logger.error("HTTP error in CountTokens handler: %s", e) + verbose_logger.debug("HTTP error in CountTokens handler: %s", e) raise BedrockError( status_code=e.response.status_code, message=e.response.text, diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py new file mode 100644 index 00000000000..23c3dc54c95 --- /dev/null +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -0,0 +1,99 @@ +from typing import Final + +import httpx +from pydantic import JsonValue, TypeAdapter +from typing_extensions import NotRequired, ReadOnly, TypedDict + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.anthropic.count_tokens.transformation import AnthropicCountTokensConfig +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing +from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client + +MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens" +MANTLE_ANTHROPIC_VERSION: Final = "2023-06-01" + + +class MantleCountTokensRequest(TypedDict): + messages: ReadOnly[list[dict[str, JsonValue]]] + system: ReadOnly[NotRequired[JsonValue]] + tools: ReadOnly[NotRequired[list[dict[str, JsonValue]]]] + + +_COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest) +_COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue]) + + +class BedrockMantleCountTokensHandler(BaseAWSLLM): + """Counts tokens through Anthropic's count_tokens on the bedrock-mantle endpoint. + + Claude models that Bedrock offers only through cross-region inference answer 400 on + bedrock-runtime's CountTokens; AWS documents Mantle's /anthropic/v1/messages/count_tokens + as the way to count them, with the base model id and the deployment's AWS credentials + """ + + def __init__(self, anthropic_config: AnthropicCountTokensConfig | None = None) -> None: + super().__init__() + self._anthropic_config: Final = anthropic_config or AnthropicCountTokensConfig() + + def get_mantle_count_tokens_endpoint(self, aws_region_name: str) -> str: + messages_url: Final = build_mantle_messages_url( + api_base=None, aws_bedrock_runtime_endpoint=None, region=aws_region_name + ) + return f"{messages_url}{MANTLE_COUNT_TOKENS_SUFFIX}" + + async def handle_count_tokens_request( + self, + request_data: dict[str, object], + litellm_params: dict[str, object], + resolved_model: str, + client: AsyncHTTPHandler | None = None, + ) -> dict[str, JsonValue]: + try: + request: Final = _COUNT_REQUEST.validate_python(request_data) + aws_region_name: Final = self._get_aws_region_name( + optional_params=litellm_params, model=resolved_model, model_id=None + ) + endpoint_url: Final = self.get_mantle_count_tokens_endpoint(aws_region_name) + body: Final = self._anthropic_config.transform_request_to_count_tokens( + model=resolved_model, + messages=request["messages"], + tools=request.get("tools"), + system=request.get("system"), + ) + verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url) + api_key: Final = litellm_params.get("api_key") + signed_headers, signed_body = await run_aws_signing( + self._sign_request, + service_name="bedrock", + headers={"Content-Type": "application/json", "anthropic-version": MANTLE_ANTHROPIC_VERSION}, + optional_params=litellm_params, + request_data=body, + api_base=endpoint_url, + model=resolved_model, + api_key=api_key if isinstance(api_key, str) else None, + ) + async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) + response: Final = await async_client.post( + endpoint_url, headers=signed_headers, data=signed_body, timeout=30.0 + ) + if response.status_code != 200: + raise BedrockError( + status_code=response.status_code, + message=response.text, + headers=response.headers, + response=response, + ) + return _COUNT_RESPONSE.validate_json(response.content) + except BedrockError: + raise + except httpx.HTTPStatusError as e: + raise BedrockError( + status_code=e.response.status_code, + message=e.response.text, + headers=e.response.headers, + response=e.response, + ) + except Exception as e: + raise BedrockError(status_code=500, message=f"bedrock-mantle count_tokens processing error: {e}") diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 9d39430cfa9..4b115686773 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -188,26 +188,12 @@ async def google_count_tokens(request: Request, model_name: str): request=token_request, call_endpoint=True, ) - if token_response is not None: - # cast the response to the well known format - original_response: Final[dict] = token_response.original_response or {} - if original_response: - return TokenCountDetailsResponse( - totalTokens=original_response.get("totalTokens", 0), - promptTokensDetails=original_response.get("promptTokensDetails", []), - ) - else: - return TokenCountDetailsResponse( - totalTokens=token_response.total_tokens or 0, - promptTokensDetails=[], - ) - - ######################################################### - # Return the response in the well known format - ######################################################### + if token_response is None: + return TokenCountDetailsResponse(totalTokens=0, promptTokensDetails=[]) + original_response: Final[dict] = token_response.original_response or {} return TokenCountDetailsResponse( - totalTokens=0, - promptTokensDetails=[], + totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0, + promptTokensDetails=original_response.get("promptTokensDetails", []), ) diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py new file mode 100644 index 00000000000..f9d4f7f91bb --- /dev/null +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -0,0 +1,1072 @@ +import asyncio +import json +import re +import signal +import socket +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping, Sequence +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import ExitStack +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import psutil +import pytest +import yaml +from anthropic.types import MessageCountTokensToolParam, MessageParam +from integration._support.bedrock_runtime_peer import answer, marker_of, target_of +from integration._support.bedrock_runtime_peer import respond as runtime_generation +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_OPUS: Final = "global.anthropic.claude-opus-4-8" +_OPUS_BASE: Final = "anthropic.claude-opus-4-8" +_SONNET: Final = "anthropic.claude-sonnet-4-6" +_NOVA: Final = "amazon.nova-lite-v1:0" +_REGION: Final = "us-east-1" +_OWNED_OPUS: Final = "mantle-opus" +_OWNED_SONNET: Final = "mantle-sonnet" +_OWNED_NOVA: Final = "mantle-nova" +_OWNED_MODELS: Final = MappingProxyType({_OWNED_OPUS: _OPUS, _OWNED_SONNET: _SONNET, _OWNED_NOVA: _NOVA}) +_ACCESS_KEY: Final = "AKIAINTEGRATIONMANTLE" +_SECRET_KEY: Final = "integration-mantle-secret" +_SIGV4_SCOPE: Final = f"/{_REGION}/bedrock/aws4_request" +_MANTLE_TARGET: Final = "/anthropic/v1/messages/count_tokens" +_MANTLE_VERSION: Final = "2023-06-01" +_MANTLE_COUNT: Final = 4242 +_RUNTIME_COUNT: Final = 1345 +_UNSUPPORTED: Final = "The provided model doesn't support counting tokens." +_REJECTION: Final = "scripted mantle rejection" +_COUNT_TARGET: Final = re.compile(r"^/model/(.+)/count-tokens$") +_INVOKE_TARGET: Final = re.compile(r"^/model/(.+)/invoke$") +_SCRIPTED_STATUS: Final = re.compile(r"status=(\d{3})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) + +_SDK_MESSAGES: Final[list[MessageParam]] = [{"role": "user", "content": "Count this message"}] +_MESSAGES: Final = _JSON_LIST.validate_json(json.dumps(_SDK_MESSAGES)) +_SYSTEM: Final = "You are a terse assistant that answers in one sentence" +_SYSTEM_BLOCKS: Final[list[JsonValue]] = [ + {"type": "text", "text": "You are a terse assistant"}, + {"type": "text", "text": "Answer in one sentence"}, +] +_SDK_TOOLS: Final[list[MessageCountTokensToolParam]] = [ + { + "name": "get_weather", + "description": "Look up the current weather for a city", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string", "description": "City to look up"}}, + "required": ["city"], + }, + } +] +_TOOLS: Final = _JSON_LIST.validate_json(json.dumps(_SDK_TOOLS)) +_GEMINI_BODY: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "Count this"}]}]} +_GEMINI_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this"}] + + +def _mantle_body(**fields: JsonValue) -> dict[str, JsonValue]: + return {"model": _OPUS_BASE, "messages": _MESSAGES, **fields} + + +_MANTLE_BARE: Final = _mantle_body() +_MANTLE_FULL: Final = _mantle_body(system=_SYSTEM, tools=_TOOLS) + + +def _json_reply(status: int, payload: Mapping[str, JsonValue]) -> Reply: + return Reply(status=status, body=json.dumps(payload).encode()) + + +def _runtime_count(request: Request, model: str) -> Reply: + scripted: Final = _SCRIPTED_STATUS.search(request.body.decode(errors="replace")) + if scripted is not None: + return _json_reply(int(scripted.group(1)), {"message": f"scripted {scripted.group(1)}"}) + if "sonnet" in model: + return _json_reply(200, {"inputTokens": _RUNTIME_COUNT}) + return _json_reply(400, {"message": _UNSUPPORTED}) + + +def _invoke_reply(marker: str) -> Reply: + return _json_reply( + 200, + { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "model": _OPUS_BASE, + "content": [{"type": "text", "text": answer(marker)}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, + }, + ) + + +def _runtime(request: Request) -> Reply: + target: Final = target_of(request) + counted: Final = _COUNT_TARGET.match(target) + if counted is not None: + return _runtime_count(request, counted.group(1)) + if _INVOKE_TARGET.match(target): + return _invoke_reply(marker_of(request)) + return runtime_generation(request) + + +def _mantle_counted(_request: Request) -> Reply: + return _json_reply(200, {"input_tokens": _MANTLE_COUNT}) + + +def _rejected(status: int) -> Reply: + return _json_reply(status, {"type": "error", "error": {"type": "invalid_request_error", "message": _REJECTION}}) + + +def _rejecting(status: int) -> Callable[[Request], Reply]: + def count(_request: Request) -> Reply: + return _rejected(status) + + return count + + +def _anthropic_message(message: JsonValue) -> bool: + return isinstance(message, dict) and message.get("role") in ("user", "assistant") + + +def _anthropic_tool(tool: JsonValue) -> bool: + return isinstance(tool, dict) and isinstance(tool.get("name"), str) and isinstance(tool.get("input_schema"), dict) + + +def _strict(request: Request) -> Reply: + body: Final = _JSON_OBJECT.validate_json(request.body) + messages: Final = body.get("messages") + tools: Final = body.get("tools", []) + accepted: Final = ( + isinstance(messages, list) + and all(map(_anthropic_message, messages)) + and isinstance(body.get("system", ""), (str, list)) + and isinstance(tools, list) + and all(map(_anthropic_tool, tools)) + ) + return _mantle_counted(request) if accepted else _rejected(400) + + +def _mantle(count: Callable[[Request], Reply] = _mantle_counted) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == _MANTLE_TARGET: + return count(request) + return _json_reply(404, {"error": f"unscripted mantle target {request.target}"}) + + return respond + + +def _mantle_environment(port: int) -> Mapping[str, str]: + return {"BEDROCK_MANTLE_API_BASE": f"http://127.0.0.1:{port}"} + + +_INHERITED_BEARER: Final = ("AWS_BEARER_TOKEN_BEDROCK",) + + +def _reserved_port() -> int: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + return reserve.getsockname()[1] + + +def _closed_port_url() -> str: + return f"http://127.0.0.1:{_reserved_port()}" + + +@pytest.fixture(scope="module") +def mantle_port() -> int: + return _reserved_port() + + +@pytest.fixture(scope="module") +def counting_proxy(tmp_path_factory: pytest.TempPathFactory, mantle_port: int) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("mantle-count"), + _mantle_environment(mantle_port), + workers=1, + remove_environment=_INHERITED_BEARER, + ) as owned, + ): + yield owned + + +def _litellm_params(model: str, api_base: str) -> dict[str, JsonValue]: + return { + "model": f"bedrock/{model}", + "api_base": api_base, + "aws_access_key_id": _ACCESS_KEY, + "aws_secret_access_key": _SECRET_KEY, + "aws_region_name": _REGION, + } + + +def _owned_config(path: Path, runtime_url: str, settings: Mapping[str, JsonValue]) -> Path: + config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + litellm_settings: Final = object_value(config["litellm_settings"]) + path.write_text( + yaml.safe_dump( + { + **config, + "model_list": [ + {"model_name": name, "litellm_params": _litellm_params(model, runtime_url)} + for name, model in _OWNED_MODELS.items() + ], + "litellm_settings": {**litellm_settings, **settings}, + } + ) + ) + return path + + +def _deployment(scenario: Scenario, api_base: str, model: str = _OPUS) -> str: + return scenario.model(model_info=None, api_key=None, **_litellm_params(model, api_base)) + + +def _bare(model: str) -> dict[str, JsonValue]: + return {"model": model, "messages": _MESSAGES} + + +def _full(model: str) -> dict[str, JsonValue]: + return {**_bare(model), "system": _SYSTEM, "tools": _TOOLS} + + +def _count(gateway: Gateway, body: Mapping[str, JsonValue]) -> httpx.Response: + return gateway.request("POST", "/v1/messages/count_tokens", body) + + +def _payload(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text + return _JSON_OBJECT.validate_json(response.content) + + +def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int: + response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "false"}) + payload: Final = _payload(response) + total: Final = payload["total_tokens"] + assert payload["tokenizer_type"] not in ("bedrock_api", "bedrock_mantle_api"), response.text + assert isinstance(total, int) and total > 0 and total not in (_MANTLE_COUNT, _RUNTIME_COUNT), response.text + return total + + +def _assert_sigv4(request: Request) -> None: + authorization: Final = request.headers.get("authorization", "") + assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={_ACCESS_KEY}/"), request.headers + assert _SIGV4_SCOPE in authorization, authorization + + +def _mantle_requests(requests: Sequence[Request]) -> tuple[dict[str, JsonValue], ...]: + for request in requests: + assert (request.method, request.target) == ("POST", _MANTLE_TARGET), request.target + assert request.headers["anthropic-version"] == _MANTLE_VERSION, request.headers + assert request.headers["content-type"] == "application/json", request.headers + _assert_sigv4(request) + return tuple(_JSON_OBJECT.validate_json(request.body) for request in requests) + + +def _mantle_bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]: + return _mantle_requests(wire.drain()) + + +def _runtime_count_targets(wire: Wire) -> tuple[str, ...]: + counts: Final = tuple(request for request in wire.drain() if _COUNT_TARGET.match(target_of(request))) + for request in counts: + assert request.method == "POST", request.method + _assert_sigv4(request) + return tuple(target_of(request) for request in counts) + + +def _clients(stack: ExitStack, base_url: str, count: int, timeout: float = 30) -> tuple[httpx.Client, ...]: + return tuple( + stack.enter_context(httpx.Client(base_url=base_url, timeout=timeout, trust_env=False)) for _ in range(count) + ) + + +def _counted_on(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue]: + response: Final = client.post( + "/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {key}"} + ) + return response.status_code, _JSON_OBJECT.validate_json(response.content).get("input_tokens") + + +def _generated_then_counted( + client: httpx.Client, key: str, model: str, body: Mapping[str, JsonValue] +) -> tuple[int, int, JsonValue]: + generated: Final = client.post( + "/v1/messages", + json={ + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"Generate before counting marker-{uuid.uuid4().hex}"}], + }, + headers={"Authorization": f"Bearer {key}"}, + ) + return generated.status_code, *_counted_on(client, key, body) + + +def _counted_or_dropped(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue] | None: + try: + return _counted_on(client, key, body) + except httpx.TransportError: + return None + + +def _probed_then_counted_or_dropped( + client: httpx.Client, key: str, body: Mapping[str, JsonValue], probed: SimpleQueue[int] +) -> tuple[int, tuple[int, JsonValue] | None]: + port: Final = _local_port(client) + probed.put(port) + return port, _counted_or_dropped(client, key, body) + + +def _local_port(client: httpx.Client) -> int: + with client.stream("GET", "/health/liveliness") as response: + port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1]) + response.read() + assert response.status_code == 200, response.text + return port + + +def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]: + return frozenset( + connection.raddr.port + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.raddr and connection.laddr.port == proxy_port + ) + + +def _holding(held: SimpleQueue[str], release: threading.Event, seconds: float) -> Callable[[Request], Reply]: + def hold(request: Request) -> Reply: + held.put(request.target) + assert release.wait(timeout=seconds), "Held count was never released" + return _mantle_counted(request) + + return hold + + +def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def _async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +@pytest.mark.parametrize("system", [_SYSTEM, _SYSTEM_BLOCKS], ids=["string", "blocks"]) +def test_messages_count_tokens_counts_through_mantle_when_the_runtime_cannot( + counting_proxy: OwnedProxy, mantle_port: int, system: JsonValue +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) + assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) + + +def test_messages_count_tokens_without_system_or_tools_sends_a_bare_body_to_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, _bare(model)) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_BARE,) + + +def test_utils_token_counter_call_endpoint_reports_the_mantle_tokenizer( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} + ) + payload: Final = _payload(response) + assert (payload["total_tokens"], payload["tokenizer_type"]) == (_MANTLE_COUNT, "bedrock_mantle_api") + assert payload["original_response"] == {"input_tokens": _MANTLE_COUNT}, response.text + assert (payload["request_model"], payload["model_used"]) == (model, _OPUS), response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_FULL,) + + +def test_gemini_count_tokens_route_reports_the_mantle_total(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) + assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == ({"model": _OPUS_BASE, "messages": _GEMINI_MESSAGES},) + + +def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock_counts( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url, _SONNET) + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) + assert _payload(response) == {"totalTokens": _RUNTIME_COUNT, "promptTokensDetails": []}, response.text + assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) + assert mantle.drain() == () + + +def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"} + ) + assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_BARE,) + + +def test_anthropic_sdk_count_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + counted: Final = _anthropic_client(gateway).messages.count_tokens( + model=model, messages=_SDK_MESSAGES, system=_SYSTEM, tools=_SDK_TOOLS + ) + assert counted.input_tokens == _MANTLE_COUNT, counted + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_FULL,) + + +def test_async_anthropic_sdk_count_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + counted: Final = asyncio.run( + _async_anthropic_client(gateway).messages.count_tokens(model=model, messages=_SDK_MESSAGES, system=_SYSTEM) + ) + assert counted.input_tokens == _MANTLE_COUNT, counted + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(system=_SYSTEM),) + + +def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url, _SONNET) + assert _payload(_count(gateway, _full(model))) == {"input_tokens": _RUNTIME_COUNT} + detailed: Final = gateway.request( + "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} + ) + payload: Final = _payload(detailed) + assert (payload["total_tokens"], payload["tokenizer_type"]) == (_RUNTIME_COUNT, "bedrock_api"), detailed.text + assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) * 2 + assert mantle.drain() == () + + +def test_non_claude_model_the_runtime_cannot_count_falls_back_locally_without_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + 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 _runtime_count_targets(runtime) == (f"/model/{_NOVA}/count-tokens",) + assert mantle.drain() == () + + +def test_runtime_403_on_a_claude_model_never_reaches_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + body: Final[dict[str, JsonValue]] = { + "model": model, + "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 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: + 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 mantle.drain() == () + + +def test_bedrock_passthrough_count_tokens_still_answers_the_runtime_rejection( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request("POST", "/bedrock/v1/messages/count_tokens", _full(model)) + assert response.status_code == 400, response.text + assert _UNSUPPORTED in response.text and "input_tokens" not in response.text, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert mantle.drain() == () + + +def test_responses_input_tokens_with_instructions_still_counts_locally( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": model, "input": "Count this message", "instructions": "Be terse"}, + ) + payload: Final = _payload(response) + assert len(_runtime_count_targets(runtime)) == 1 + (sent,) = _mantle_bodies(mantle) + messages: Final = sent["messages"] + assert isinstance(messages, list) and messages[0] == {"role": "system", "content": "Be terse"}, sent + assert "system" not in sent, sent + local: Final = _local_count(gateway, {"model": model, "messages": messages}) + assert payload == {"object": "response.input_tokens", "input_tokens": local}, response.text + + +@pytest.mark.parametrize("status", [400, 403, 404, 500, 503]) +def test_messages_count_tokens_falls_back_locally_when_mantle_errors( + counting_proxy: OwnedProxy, mantle_port: int, status: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_rejecting(status)), port=mantle_port) as mantle, + 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 len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_FULL,) + + +@pytest.mark.parametrize( + "body", + [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( + counting_proxy: OwnedProxy, mantle_port: int, body: bytes +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(lambda _request: Reply(body=body)), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + 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 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: + 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 len(_runtime_count_targets(runtime)) == 1 + + +def test_messages_count_tokens_falls_back_locally_when_mantle_never_answers( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + + def hold(request: Request) -> Reply: + held.put(request.target) + release.wait(timeout=120) + return Reply(drop_connection=True) + + with ExitStack() as stack: + (client,) = _clients(stack, str(gateway.client.base_url), 1, timeout=120) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=1)) + runtime: Final = stack.enter_context(wire_server(_runtime)) + mantle: Final = stack.enter_context(wire_server(_mantle(hold), port=mantle_port)) + stack.callback(release.set) + 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 not future.done() + assert future.result(timeout=120) == (200, local) + 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( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + 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 len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(system=5),) + + +@pytest.mark.parametrize( + "tools", + [5, "", "x" * 5120, ["get_weather"]], + ids=["int", "empty_string", "5kb_string", "list_of_strings"], +) +def test_messages_count_tokens_rejects_malformed_tools_without_calling_either_peer( + counting_proxy: OwnedProxy, mantle_port: int, tools: JsonValue +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + refused: Final = _count(gateway, {**_bare(model), "tools": tools}) + assert 400 <= refused.status_code < 600, refused.text + assert "input_tokens" not in refused.text, refused.text + assert runtime.drain() == () and mantle.drain() == () + assert _generated_then_counted(gateway.client, gateway.key, model, _bare(model)) == (200, 200, _MANTLE_COUNT) + assert _mantle_bodies(mantle) == (_MANTLE_BARE,) + + +@pytest.mark.parametrize("fields", [{}, {"messages": []}], ids=["missing", "empty"]) +def test_messages_count_tokens_without_messages_is_rejected( + counting_proxy: OwnedProxy, mantle_port: int, fields: dict[str, JsonValue] +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {"model": model, "system": _SYSTEM, "tools": _TOOLS, **fields}) + assert response.status_code == 400, response.text + assert "messages parameter is required" in response.text, response.text + assert runtime.drain() == () and mantle.drain() == () + + +def test_messages_count_tokens_unauthenticated_request_never_reaches_either_peer( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", _full(model), key="sk-not-a-key-this-proxy-issued" + ) + assert response.status_code == 401, response.text + assert "input_tokens" not in response.text, response.text + assert runtime.drain() == () and mantle.drain() == () + + +def test_messages_count_tokens_forwards_a_5kb_system_verbatim_to_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + system: Final = "Answer in one sentence. " * 214 + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {**_bare(model), "system": system}) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(system=system),) + + +def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value_each_to_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + fields: Final = f'"system": {json.dumps(_SYSTEM)}, "tools": {json.dumps(_TOOLS)}' + response: Final = gateway.client.post( + "/v1/messages/count_tokens", + content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}', + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + (sent,) = mantle.drain() + assert _mantle_requests((sent,)) == (_MANTLE_FULL,) + assert (sent.body.count(b'"system"'), sent.body.count(b'"tools"')) == (1, 1), sent.body + + +def test_messages_count_tokens_leaves_empty_tools_and_system_out_of_the_mantle_body( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {**_bare(model), "tools": [], "system": ""}) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_BARE,) + + +def test_messages_count_tokens_repeated_request_reaches_both_peers_each_time( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2)) + assert answers == ({"input_tokens": _MANTLE_COUNT},) * 2 + assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 + assert _mantle_bodies(mantle) == (_MANTLE_BARE,) * 2 + + +def test_disabled_token_counter_surfaces_the_mantle_error_instead_of_counting_locally( + gateway: Gateway, mantle_port: int, tmp_path: Path +) -> None: + with ExitStack() as stack: + runtime: Final = stack.enter_context(wire_server(_runtime)) + config: Final = _owned_config( + tmp_path / "disabled-token-counter.yaml", runtime.url, {"disable_token_counter": True} + ) + owned: Final = stack.enter_context( + owned_proxy_process( + gateway, + tmp_path, + _mantle_environment(mantle_port), + config=config, + workers=2, + remove_environment=_INHERITED_BEARER, + ) + ) + with wire_server(_mantle(_rejecting(403)), port=mantle_port) as refusing: + refused: Final = _count(owned.gateway, _full(_OWNED_OPUS)) + assert refused.status_code == 403, refused.text + assert _REJECTION in refused.text and "input_tokens" not in refused.text, refused.text + assert _mantle_bodies(refusing) == (_MANTLE_FULL,) + with wire_server(_mantle(), port=mantle_port) as counting: + assert _payload(_count(owned.gateway, _full(_OWNED_OPUS))) == {"input_tokens": _MANTLE_COUNT} + assert _mantle_bodies(counting) == (_MANTLE_FULL,) + assert _payload(_count(owned.gateway, _full(_OWNED_SONNET))) == {"input_tokens": _RUNTIME_COUNT} + rejected: Final = _count(owned.gateway, _full(_OWNED_NOVA)) + assert rejected.status_code == 400, rejected.text + assert _UNSUPPORTED in rejected.text and "input_tokens" not in rejected.text, rejected.text + assert counting.drain() == () + assert len(_runtime_count_targets(runtime)) == 4 + + +def test_mantle_outage_between_concurrent_waves_falls_back_then_recovers( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ExitStack() as stack: + clients: Final = _clients(stack, str(gateway.client.base_url), 8) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + runtime: Final = stack.enter_context(wire_server(_runtime)) + 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) + + def count_only(client: httpx.Client) -> tuple[int, JsonValue]: + return _counted_on(client, gateway.key, body) + + with wire_server(_mantle(), port=mantle_port) as mantle: + 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) + 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) + assert len(_runtime_count_targets(runtime)) == 3 * len(clients) + + +def test_slow_mantle_holds_concurrent_counts_without_stalling_the_proxy( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + with ExitStack() as stack: + clients: Final = _clients(stack, str(gateway.client.base_url), 6) + runtime: Final = stack.enter_context(wire_server(_runtime)) + mantle: Final = stack.enter_context(wire_server(_mantle(_holding(held, release, 20)), port=mantle_port)) + scenario: Final = stack.enter_context(gateway.scenario()) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + stack.callback(release.set) + model: Final = _deployment(scenario, runtime.url) + futures: Final = tuple( + pool.submit(_generated_then_counted, client, gateway.key, model, _full(model)) for client in clients + ) + eventually(held.qsize, lambda size: size == len(clients), seconds=30) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + assert _local_count(gateway, _full(model)) > 0 + assert not any(future.done() for future in futures) + release.set() + assert tuple(future.result(timeout=30) for future in futures) == ((200, 200, _MANTLE_COUNT),) * len(clients) + assert _mantle_bodies(mantle) == (_MANTLE_FULL,) * len(clients) + assert len(_runtime_count_targets(runtime)) == len(clients) + + +def test_worker_sigkill_mid_burst_leaves_the_sibling_counting( + gateway: Gateway, mantle_port: int, tmp_path: Path +) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + with ExitStack() as stack: + runtime: Final = stack.enter_context(wire_server(_runtime)) + mantle: Final = stack.enter_context(wire_server(_mantle(_holding(held, release, 60)), port=mantle_port)) + config: Final = _owned_config(tmp_path / "worker-kill.yaml", runtime.url, {}) + owned: Final = stack.enter_context( + owned_proxy_process( + gateway, + tmp_path, + _mantle_environment(mantle_port), + config=config, + workers=2, + remove_environment=_INHERITED_BEARER, + ) + ) + body: Final = _full(_OWNED_OPUS) + proxy_url: Final = owned.gateway.client.base_url + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + clients: Final = _clients(stack, str(proxy_url), 12) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + stack.callback(release.set) + probed: Final[SimpleQueue[int]] = SimpleQueue() + futures: Final[tuple[Future[tuple[int, tuple[int, JsonValue] | None]], ...]] = tuple( + pool.submit(_probed_then_counted_or_dropped, client, owned.gateway.key, body, probed) for client in clients + ) + eventually(held.qsize, lambda size: size == len(clients), seconds=30) + ports: Final = frozenset(probed.get_nowait() for _ in clients) + shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & ports for pid in workers} + assert sum(map(len, shares.values())) == len(clients), shares + victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) + psutil.Process(victim).send_signal(signal.SIGKILL) + release.set() + for port, result in (future.result(timeout=60) for future in futures): + assert result == (None if port in shares[victim] else (200, _MANTLE_COUNT)), (port, result, shares) + second_wave: Final = _clients(stack, str(proxy_url), 6) + assert tuple(_counted_on(client, owned.gateway.key, body) for client in second_wave) == ( + (200, _MANTLE_COUNT), + ) * len(second_wave) + assert len(_mantle_bodies(mantle)) == len(clients) + len(second_wave) + eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120) + assert owned.process.poll() is None + + +def _streamed_text(text: str) -> str: + events: Final = tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in text.splitlines() + if line.startswith("data: {") + ) + return "".join(_delta_content(event) for event in events) + + +def _delta_content(event: dict[str, JsonValue]) -> str: + choices: Final = event.get("choices") + if not isinstance(choices, list) or not choices: + return "" + delta: Final = object_value(choices[0]).get("delta") + if not isinstance(delta, dict): + return "" + content: Final = delta.get("content") + return content if isinstance(content, str) else "" + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_chat_completions_on_the_same_deployment_still_generate( + counting_proxy: OwnedProxy, mantle_port: int, stream: bool +) -> None: + gateway: Final = counting_proxy.gateway + marker: Final = uuid.uuid4().hex + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"chat control marker-{marker}"}], + "stream": stream, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + generated: Final = _streamed_text(response.text) if stream else response.text + assert answer(marker) in generated, response.text + assert not stream or response.text.rstrip().endswith("data: [DONE]"), response.text + (sent,) = runtime.drain() + assert target_of(sent) == f"/model/{_OPUS}/{'converse-stream' if stream else 'converse'}", sent.target + assert mantle.drain() == () + + +def test_messages_endpoint_on_the_same_deployment_still_generates(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + marker: Final = uuid.uuid4().hex + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": f"marker-{marker}"}]}, + ) + assert response.status_code == 200, response.text + assert answer(marker) in response.text, response.text + (sent,) = runtime.drain() + assert target_of(sent) == f"/model/{_OPUS}/invoke", sent.target + assert mantle.drain() == () + + +def test_responses_endpoint_on_the_same_deployment_still_generates( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + marker: Final = uuid.uuid4().hex + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": f"marker-{marker}"}) + assert response.status_code == 200, response.text + assert answer(marker) in response.text, response.text + (sent,) = runtime.drain() + assert target_of(sent) == f"/model/{_OPUS}/converse", sent.target + assert mantle.drain() == () + + +def test_utils_token_counter_local_mode_never_calls_either_peer(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + assert _local_count(gateway, _full(model)) > 0 + assert runtime.drain() == () and mantle.drain() == () diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py index e4cd7d9dfa8..451792109de 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -3,6 +3,7 @@ Test to verify the Google GenAI proxy API endpoints """ +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -179,3 +180,30 @@ def test_google_count_tokens_unchanged(): assert response.status_code == 200 body = response.json() assert body["totalTokens"] == 7 + + +def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs(): + """A provider that counts in Anthropic shape (Bedrock, Anthropic) carries no totalTokens key, so the normalized total must win over the raw response's missing field.""" + try: + client: Final = _build_test_client() + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + fake_response: Final = MagicMock() + fake_response.original_response = {"input_tokens": 2167} + fake_response.total_tokens = 2167 + + with patch( + "litellm.proxy.proxy_server.token_counter", + new_callable=AsyncMock, + return_value=fake_response, + ): + response: Final = client.post( + "/v1beta/models/claude-opus-4-8:countTokens", + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + ) + + assert response.status_code == 200 + body: Final = response.json() + assert body["totalTokens"] == 2167 + assert body["promptTokensDetails"] == [] diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py new file mode 100644 index 00000000000..71b451f834c --- /dev/null +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py @@ -0,0 +1,109 @@ +import json +from collections.abc import Mapping +from typing import Final + +import httpx +import pytest + +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + +MANTLE_COUNT_URL: Final = "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages/count_tokens" +LITELLM_PARAMS: Final = { + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", +} +REQUEST: Final[dict[str, object]] = { + "model": "global.anthropic.claude-opus-4-8", + "messages": [{"role": "user", "content": "The quick brown fox"}], + "system": "You are a terse assistant.", + "tools": [{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}], +} + + +class _MantleEndpoint: + def __init__(self, status_code: int, body: Mapping[str, object]) -> None: + self.status_code: Final = status_code + self.body: Final = body + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + return httpx.Response(self.status_code, json=dict(self.body), request=request) + + def only_request(self) -> httpx.Request: + assert len(self.requests) == 1, self.requests + return self.requests[0] + + +def _client(endpoint: _MantleEndpoint) -> AsyncHTTPHandler: + return AsyncHTTPHandler(transport=httpx.MockTransport(endpoint)) + + +@pytest.fixture(autouse=True) +def _sigv4_only(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + + +@pytest.mark.asyncio +async def test_counts_on_mantle_with_the_base_model_and_the_deployment_credentials() -> None: + mantle: Final = _MantleEndpoint(200, {"input_tokens": 2177}) + + result: Final = await BedrockMantleCountTokensHandler().handle_count_tokens_request( + request_data=dict(REQUEST), + litellm_params=dict(LITELLM_PARAMS), + resolved_model="anthropic.claude-opus-4-8", + client=_client(mantle), + ) + + assert result == {"input_tokens": 2177} + posted: Final = mantle.only_request() + assert str(posted.url) == MANTLE_COUNT_URL + assert posted.headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "us-east-1/bedrock/aws4_request" in posted.headers["Authorization"] + assert posted.headers["anthropic-version"] == "2023-06-01" + assert json.loads(posted.content) == { + "model": "anthropic.claude-opus-4-8", + "messages": REQUEST["messages"], + "system": REQUEST["system"], + "tools": REQUEST["tools"], + } + + +@pytest.mark.asyncio +async def test_bedrock_mantle_api_base_env_names_the_host(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://vpce-abc.bedrock-mantle.us-east-1.vpce.api.aws") + mantle: Final = _MantleEndpoint(200, {"input_tokens": 3}) + + await BedrockMantleCountTokensHandler().handle_count_tokens_request( + request_data=dict(REQUEST), + litellm_params=dict(LITELLM_PARAMS), + resolved_model="anthropic.claude-opus-4-8", + client=_client(mantle), + ) + + assert ( + str(mantle.only_request().url) + == "https://vpce-abc.bedrock-mantle.us-east-1.vpce.api.aws/anthropic/v1/messages/count_tokens" + ) + + +@pytest.mark.asyncio +async def test_non_200_answers_raise_bedrock_error_with_mantle_status() -> None: + mantle: Final = _MantleEndpoint( + 404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}} + ) + + with pytest.raises(BedrockError) as raised: + await BedrockMantleCountTokensHandler().handle_count_tokens_request( + request_data=dict(REQUEST), + litellm_params=dict(LITELLM_PARAMS), + resolved_model="anthropic.claude-sonnet-5-5", + client=_client(mantle), + ) + + assert raised.value.status_code == 404 + assert "does not exist" in raised.value.message diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py new file mode 100644 index 00000000000..a3dd8021f73 --- /dev/null +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -0,0 +1,140 @@ +from collections.abc import Mapping +from typing import Final + +import httpx +import pytest +from pydantic import JsonValue, TypeAdapter + +from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + +RUNTIME_HOST: Final = "bedrock-runtime.us-east-1.amazonaws.com" +MANTLE_HOST: Final = "bedrock-mantle.us-east-1.api.aws" +UNSUPPORTED: Final = {"message": "The provided model doesn't support counting tokens."} +DEPLOYMENT: Final = { + "litellm_params": { + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + } +} +MESSAGES: Final = [{"role": "user", "content": "The quick brown fox jumps over the lazy dog."}] +_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue]) + + +class _Bedrock: + def __init__(self, runtime: tuple[int, Mapping[str, object]], mantle: tuple[int, Mapping[str, object]]) -> None: + self.runtime: Final = runtime + self.mantle: Final = mantle + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + status_code, body = self.runtime if request.url.host == RUNTIME_HOST else self.mantle + return httpx.Response(status_code, json=dict(body), request=request) + + def posted_hosts(self) -> tuple[str, ...]: + return tuple(request.url.host for request in self.requests) + + +def _counter(bedrock: _Bedrock) -> BedrockTokenCounter: + return BedrockTokenCounter(client=AsyncHTTPHandler(transport=httpx.MockTransport(bedrock))) + + +@pytest.fixture(autouse=True) +def _sigv4_only(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + + +@pytest.mark.asyncio +async def test_claude_model_bedrock_runtime_cannot_count_is_counted_on_mantle() -> None: + bedrock: Final = _Bedrock(runtime=(400, UNSUPPORTED), mantle=(200, {"input_tokens": 2177})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + system="You are a terse assistant.", + ) + + assert result is not None + assert result.error is False + assert result.total_tokens == 2177 + assert result.tokenizer_type == "bedrock_mantle_api" + assert result.original_response == {"input_tokens": 2177} + assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) + mantle_body: Final = _JSON_BODY.validate_json(bedrock.requests[1].content) + assert mantle_body["model"] == "anthropic.claude-opus-4-8" + assert mantle_body["system"] == "You are a terse assistant." + + +@pytest.mark.asyncio +async def test_bedrock_runtime_count_is_kept_when_it_answers() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-sonnet-4-6", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-sonnet-4-6", + ) + + assert result is not None + assert result.total_tokens == 1353 + assert result.tokenizer_type == "bedrock_api" + assert bedrock.posted_hosts() == (RUNTIME_HOST,) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_to_use", "runtime"), + ( + ("global.anthropic.claude-opus-4-8", (403, {"Message": "not authorized to perform: bedrock:CountTokens"})), + ("amazon.nova-pro-v1:0", (400, UNSUPPORTED)), + ), +) +async def test_other_bedrock_runtime_failures_are_not_retried_on_mantle( + model_to_use: str, runtime: tuple[int, Mapping[str, object]] +) -> None: + bedrock: Final = _Bedrock(runtime=runtime, mantle=(200, {"input_tokens": 2177})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use=model_to_use, + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model=model_to_use, + ) + + assert result is not None + assert result.error is True + assert result.status_code == runtime[0] + assert result.tokenizer_type == "bedrock_api" + assert bedrock.posted_hosts() == (RUNTIME_HOST,) + + +@pytest.mark.asyncio +async def test_mantle_failure_is_reported_with_its_status() -> None: + bedrock: Final = _Bedrock( + runtime=(400, UNSUPPORTED), + mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}), + ) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-sonnet-5-5", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-sonnet-5-5", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 404 + assert result.tokenizer_type == "bedrock_mantle_api" + assert "does not exist" in (result.error_message or "") + assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index fcd72e3777c..4bd48e1effc 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from dotenv import load_dotenv +from pydantic import JsonValue load_dotenv() @@ -20,6 +21,7 @@ from litellm import Router from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.anthropic_endpoints.endpoints import ( count_tokens as anthropic_count_tokens, @@ -753,41 +755,45 @@ def test_vertex_ai_partner_models_token_counting_endpoint(vertex_location): ) +class _FailingRuntimeHandler(BedrockCountTokensHandler): + def __init__(self, error: Exception) -> None: + super().__init__() + self._error = error + + async def handle_count_tokens_request( + self, + request_data: dict[str, object], + litellm_params: dict[str, object], + resolved_model: str, + client: AsyncHTTPHandler | None = None, + ) -> dict[str, JsonValue]: + raise self._error + + @pytest.mark.asyncio async def test_bedrock_token_counter_error_propagation_bedrock_error(): """ Test that BedrockTokenCounter properly returns error response when BedrockError is raised. Verifies that the status code and error message are preserved. """ - counter = BedrockTokenCounter() + counter = BedrockTokenCounter( + runtime_handler=_FailingRuntimeHandler(BedrockError(status_code=429, message="Rate limit exceeded")) + ) - # Mock the handler to raise BedrockError with specific status code - with patch.object( - counter, "count_tokens", wraps=counter.count_tokens - ) as mock_count: - # We need to patch at the handler level - with patch( - "litellm.llms.bedrock.count_tokens.bedrock_token_counter.BedrockCountTokensHandler" - ) as MockHandler: - mock_handler_instance = MockHandler.return_value - mock_handler_instance.handle_count_tokens_request = AsyncMock( - side_effect=BedrockError(status_code=429, message="Rate limit exceeded") - ) + result = await counter.count_tokens( + model_to_use="anthropic.claude-3-sonnet", + messages=[{"role": "user", "content": "hello"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="bedrock/anthropic.claude-3-sonnet", + ) - result = await counter.count_tokens( - model_to_use="anthropic.claude-3-sonnet", - messages=[{"role": "user", "content": "hello"}], - contents=None, - deployment={"litellm_params": {}}, - request_model="bedrock/anthropic.claude-3-sonnet", - ) - - assert result is not None - assert result.error is True - assert result.status_code == 429 - assert "Rate limit exceeded" in result.error_message - assert result.tokenizer_type == "bedrock_api" - assert result.total_tokens == 0 + assert result is not None + assert result.error is True + assert result.status_code == 429 + assert "Rate limit exceeded" in result.error_message + assert result.tokenizer_type == "bedrock_api" + assert result.total_tokens == 0 @pytest.mark.asyncio @@ -795,28 +801,20 @@ async def test_bedrock_token_counter_error_propagation_generic_exception(): """ Test that BedrockTokenCounter returns error response with 500 status for generic exceptions. """ - counter = BedrockTokenCounter() + counter = BedrockTokenCounter(runtime_handler=_FailingRuntimeHandler(Exception("Unexpected error"))) - with patch( - "litellm.llms.bedrock.count_tokens.bedrock_token_counter.BedrockCountTokensHandler" - ) as MockHandler: - mock_handler_instance = MockHandler.return_value - mock_handler_instance.handle_count_tokens_request = AsyncMock( - side_effect=Exception("Unexpected error") - ) + result = await counter.count_tokens( + model_to_use="anthropic.claude-3-sonnet", + messages=[{"role": "user", "content": "hello"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="bedrock/anthropic.claude-3-sonnet", + ) - result = await counter.count_tokens( - model_to_use="anthropic.claude-3-sonnet", - messages=[{"role": "user", "content": "hello"}], - contents=None, - deployment={"litellm_params": {}}, - request_model="bedrock/anthropic.claude-3-sonnet", - ) - - assert result is not None - assert result.error is True - assert result.status_code == 500 - assert "Unexpected error" in result.error_message + assert result is not None + assert result.error is True + assert result.status_code == 500 + assert "Unexpected error" in result.error_message @pytest.mark.asyncio From 0b148c95eeb031c57a2e31e1ba75faf3fe7ee7ba Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:12:57 -0700 Subject: [PATCH 09/36] fix(token_counter): price a base64 PDF document per page instead of as one image (#45301) Backport of #45301 to stable/1.104.x. Cherry-picked from 50ab9a3211254c1b8031fa0014cdcbb79718e6c4 (main). Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/constants.py | 4 + litellm/litellm_core_utils/token_counter.py | 60 +++- tests/integration/_support/pdf_document.py | 149 +++++++++ .../test_pdf_document_count_tokens_chaos.py | 226 ++++++++++++++ .../test_pdf_document_count_tokens_wire.py | 273 ++++++++++++++++ ...df_document_pre_call_checks_owned_proxy.py | 292 ++++++++++++++++++ .../litellm_core_utils/test_token_counter.py | 104 ++++++- 7 files changed, 1098 insertions(+), 10 deletions(-) create mode 100644 tests/integration/_support/pdf_document.py create mode 100644 tests/integration/providers/test_pdf_document_count_tokens_chaos.py create mode 100644 tests/integration/providers/test_pdf_document_count_tokens_wire.py create mode 100644 tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py diff --git a/litellm/constants.py b/litellm/constants.py index 3b14d99bfc3..b37a648835c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -649,6 +649,10 @@ EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE: Final = float( ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv("ANTHROPIC_TOKEN_COUNTING_BETA_VERSION", "token-counting-2024-11-01") ANTHROPIC_SKILLS_API_BETA_VERSION: Final = "skills-2025-10-02" ANTHROPIC_BATCHES_ROUTE: Final = "/v1/messages/batches" +ANTHROPIC_IMAGE_MAX_LONG_EDGE_PX: Final = 1568 +ANTHROPIC_IMAGE_MAX_PIXELS: Final = 1_150_000 +ANTHROPIC_IMAGE_PIXELS_PER_TOKEN: Final = 750 +PDF_DATA_URL_PREFIX: Final = "data:application/pdf;base64," VERTEX_BATCH_PREDICTION_JOBS_ROUTE: Final = "batchPredictionJobs" ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES: Final = { "low": 1, diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index cdd2d0654be..b4a0d7cca21 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -2,6 +2,7 @@ ## Helper utilities for token counting import base64 import io +import math import struct from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from typing import Final, Literal, cast @@ -16,6 +17,9 @@ import litellm from litellm import verbose_logger from litellm._lazy_imports import _get_default_encoding from litellm.constants import ( + ANTHROPIC_IMAGE_MAX_LONG_EDGE_PX, + ANTHROPIC_IMAGE_MAX_PIXELS, + ANTHROPIC_IMAGE_PIXELS_PER_TOKEN, DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_TOKEN_COUNT, DEFAULT_IMAGE_WIDTH, @@ -24,6 +28,7 @@ from litellm.constants import ( MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES, MAX_TILE_HEIGHT, MAX_TILE_WIDTH, + PDF_DATA_URL_PREFIX, TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS, TOKEN_COUNTER_MAX_CONCURRENT_COUNTS, TOKEN_COUNTER_MAX_EXACT_CHARS, @@ -771,6 +776,49 @@ def _anthropic_image_source_data( return "" +def _anthropic_rendered_page_image_tokens(width: float, height: float) -> int: + """A PDF page is rasterized within Anthropic's image limits, then billed by its pixel area.""" + if width <= 0 or height <= 0: + return 0 + area_at_max_edge: Final = ANTHROPIC_IMAGE_MAX_LONG_EDGE_PX**2 * min(width, height) / max(width, height) + return math.ceil(min(float(ANTHROPIC_IMAGE_MAX_PIXELS), area_at_max_edge) / ANTHROPIC_IMAGE_PIXELS_PER_TOKEN) + + +def _count_inline_pdf_tokens(data_url: str, count_function: TokenCounterFunction) -> int | None: + if not data_url.startswith(PDF_DATA_URL_PREFIX): + return None + try: + from pypdf import PdfReader + except ImportError: + verbose_logger.debug("pypdf is not installed, so the PDF document is priced like one image") + return None + try: + reader: Final = PdfReader(io.BytesIO(base64.b64decode(data_url[len(PDF_DATA_URL_PREFIX) :]))) + return sum( + count_function(page.extract_text() or "") + + _anthropic_rendered_page_image_tokens(float(page.mediabox.width), float(page.mediabox.height)) + for page in reader.pages + ) + except Exception as e: + verbose_logger.debug("Could not read the PDF document's pages (%s), so it is priced like one image", e) + return None + + +def _count_opaque_document_tokens( + data_url: str, + count_function: TokenCounterFunction, + use_default_image_token_count: bool, +) -> int: + pdf_tokens: Final = _count_inline_pdf_tokens(data_url, count_function) + if pdf_tokens is not None: + return pdf_tokens + return calculate_img_tokens( + data=data_url, + mode="auto", + use_default_image_token_count=use_default_image_token_count, + ) + + def _count_document_tokens( document: ChatCompletionDocumentObject | AnthropicMessagesDocumentParam, count_function: TokenCounterFunction, @@ -790,10 +838,8 @@ def _count_document_tokens( return metadata_tokens + _count_content_list( count_function, content, use_default_image_token_count, default_token_count ) - return metadata_tokens + calculate_img_tokens( - data=_anthropic_image_source_data(source), - mode="auto", - use_default_image_token_count=use_default_image_token_count, + return metadata_tokens + _count_opaque_document_tokens( + _anthropic_image_source_data(source), count_function, use_default_image_token_count ) @@ -810,11 +856,7 @@ def _count_file_tokens( name_tokens: Final = count_function(filename) if isinstance(filename, str) and filename else 0 if not isinstance(file_data, str) or not file_data: return name_tokens - return name_tokens + calculate_img_tokens( - data=file_data, - mode="auto", - use_default_image_token_count=use_default_image_token_count, - ) + return name_tokens + _count_opaque_document_tokens(file_data, count_function, use_default_image_token_count) def _count_anthropic_content( diff --git a/tests/integration/_support/pdf_document.py b/tests/integration/_support/pdf_document.py new file mode 100644 index 00000000000..5c6ad82b85f --- /dev/null +++ b/tests/integration/_support/pdf_document.py @@ -0,0 +1,149 @@ +import base64 +import json +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import accumulate +from types import MappingProxyType +from typing import Final + +from integration._support.wire import Reply +from pydantic import JsonValue + +PDF_MEDIA_TYPE: Final = "application/pdf" +COUNT_TOKENS_TARGET: Final = "/v1/messages/count_tokens" +COUNT_REFUSED: Final = Reply( + status=400, + body=json.dumps( + {"type": "error", "error": {"type": "invalid_request_error", "message": "count_tokens is not supported"}} + ).encode(), +) + + +@dataclass(frozen=True, slots=True) +class Page: + width: int = 612 + height: int = 792 + text: str | None = None + + +LETTER: Final = Page() +NARROW: Final = Page(width=100, height=1000) +_RENDERED_TOKENS: Final = MappingProxyType({(612, 792): 1534, (100, 1000): 328}) + + +def rendered_tokens(pages: Sequence[Page]) -> int: + return sum(_RENDERED_TOKENS[(page.width, page.height)] for page in pages) + + +def _escaped(text: str) -> str: + return text.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)") + + +def _content_stream(page: Page) -> bytes: + stream: Final = f"BT /F1 12 Tf 72 {page.height - 72} Td ({_escaped(page.text or '')}) Tj ET".encode() + return b"<< /Length " + str(len(stream)).encode() + b" >>\nstream\n" + stream + b"\nendstream" + + +def _page_object(page: Page, contents_id: int | None) -> bytes: + contents: Final = f" /Contents {contents_id} 0 R" if contents_id is not None else "" + return ( + f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {page.width} {page.height}]" + f" /Resources << /Font << /F1 1 0 R >> >>{contents} >>" + ).encode() + + +def _page_bodies(pages: Sequence[Page], starts: Sequence[int]) -> Iterator[bytes]: + for page, start in zip(pages, starts): + if page.text is None: + yield _page_object(page, None) + else: + yield _content_stream(page) + yield _page_object(page, start) + + +def pdf_bytes(pages: Sequence[Page]) -> bytes: + starts: Final = tuple(accumulate((1 if page.text is None else 2 for page in pages), initial=3)) + page_ids: Final = tuple(start + (0 if page.text is None else 1) for start, page in zip(starts, pages)) + kids: Final = " ".join(f"{identity} 0 R" for identity in page_ids) + bodies: Final = ( + b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>", + f"<< /Type /Pages /Kids [{kids}] /Count {len(pages)} >>".encode(), + *_page_bodies(pages, starts), + b"<< /Type /Catalog /Pages 2 0 R >>", + ) + header: Final = b"%PDF-1.4\n" + objects: Final = tuple( + f"{number} 0 obj\n".encode() + body + b"\nendobj\n" for number, body in enumerate(bodies, start=1) + ) + offsets: Final = tuple(accumulate((len(chunk) for chunk in objects), initial=len(header))) + entries: Final = b"".join(f"{offset:010d} 00000 n \n".encode() for offset in offsets[:-1]) + trailer: Final = ( + f"xref\n0 {len(bodies) + 1}\n".encode() + + b"0000000000 65535 f \n" + + entries + + f"trailer\n<< /Size {len(bodies) + 1} /Root {starts[-1]} 0 R >>\nstartxref\n{offsets[-1]}\n%%EOF\n".encode() + ) + return header + b"".join(objects) + trailer + + +def encoded(raw: bytes) -> str: + return base64.b64encode(raw).decode() + + +def data_url(media_type: str, raw: bytes) -> str: + return f"data:{media_type};base64,{encoded(raw)}" + + +def pdf_data_url(pages: Sequence[Page]) -> str: + return data_url(PDF_MEDIA_TYPE, pdf_bytes(pages)) + + +def base64_source(data: JsonValue, media_type: JsonValue = PDF_MEDIA_TYPE) -> dict[str, JsonValue]: + return {"type": "base64", "media_type": media_type, "data": data} + + +def document(source: Mapping[str, JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: + return {"type": "document", "source": dict(source), **fields} + + +def pdf_document(pages: Sequence[Page], **fields: JsonValue) -> dict[str, JsonValue]: + return document(base64_source(encoded(pdf_bytes(pages))), **fields) + + +def text_document(text: str) -> dict[str, JsonValue]: + return document({"type": "text", "media_type": "text/plain", "data": text}) + + +def chat_file(file_data: JsonValue, filename: JsonValue = "document.pdf") -> dict[str, JsonValue]: + return {"type": "file", "file": {"filename": filename, "file_data": file_data}} + + +def responses_input_file(file_data: JsonValue, filename: JsonValue = "document.pdf") -> dict[str, JsonValue]: + return {"type": "input_file", "filename": filename, "file_data": file_data} + + +def messages_body(model: str, blocks: Sequence[JsonValue], text: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": [*blocks, {"type": "text", "text": text}]}], + **fields, + } + + +def chat_body(model: str, parts: Sequence[JsonValue], text: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": [*parts, {"type": "text", "text": text}]}], + **fields, + } + + +def responses_body(model: str, items: Sequence[JsonValue], text: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "max_output_tokens": 64, + "input": [{"role": "user", "content": [*items, {"type": "input_text", "text": text}]}], + **fields, + } diff --git a/tests/integration/providers/test_pdf_document_count_tokens_chaos.py b/tests/integration/providers/test_pdf_document_count_tokens_chaos.py new file mode 100644 index 00000000000..f196018f117 --- /dev/null +++ b/tests/integration/providers/test_pdf_document_count_tokens_chaos.py @@ -0,0 +1,226 @@ +import asyncio +import os +import re +import signal +import threading +import uuid +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.pdf_document import ( + COUNT_REFUSED, + COUNT_TOKENS_TARGET, + LETTER, + messages_body, + pdf_document, + rendered_tokens, +) +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from integration.providers._cache_control_marks_support import anthropic_peer, owned_config +from pydantic import JsonValue, TypeAdapter + +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_MODEL: Final = "burst-claude" +_PROVIDER_KEY: Final = "synthetic-provider-key" +_ASK: Final = "Summarize the attached report in one sentence." +_COUNT_BURST: Final = 24 +_MESSAGE_BURST: Final = 12 +_PAGE_COUNTS: Final = (1, 3, 12) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class _Counted: + pages: int + status: int + text: str + input_tokens: int | None + + +@dataclass(frozen=True, slots=True) +class _Sent: + marker: str + status: int + text: str + call_id: str + + +def _peer(request: Request) -> Reply: + if request.target == COUNT_TOKENS_TARGET: + return COUNT_REFUSED + return anthropic_peer(request) + + +def _held(release: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert release.wait(timeout=120), "Held peer was never released" + return _peer(request) + + return respond + + +def _deployment(api_base: str) -> dict[str, JsonValue]: + return { + "model_name": _MODEL, + "litellm_params": {"model": "anthropic/claude-opus-5-5", "api_base": api_base, "api_key": _PROVIDER_KEY}, + } + + +def _count_body(pages: int) -> dict[str, JsonValue]: + blocks: Final[list[JsonValue]] = [pdf_document((LETTER,) * pages)] if pages else [] + return {"model": _MODEL, "messages": messages_body(_MODEL, blocks, _ASK)["messages"]} + + +def _input_tokens(response: httpx.Response) -> int | None: + if response.status_code != 200: + return None + counted: Final = _JSON_OBJECT.validate_json(response.content).get("input_tokens") + return counted if isinstance(counted, int) else None + + +async def _fire_counts(url: str, key: str, *, tolerate_transport_errors: bool = False) -> tuple[_Counted, ...]: + async def one(client: httpx.AsyncClient, index: int) -> _Counted: + pages: Final = _PAGE_COUNTS[index % len(_PAGE_COUNTS)] + response: Final = await client.post( + "/v1/messages/count_tokens", json=_count_body(pages), headers={"Authorization": f"Bearer {key}"} + ) + return _Counted(pages, response.status_code, response.text, _input_tokens(response)) + + async with httpx.AsyncClient(base_url=url, timeout=90, trust_env=False) as client: + results: Final = await asyncio.gather( + *(one(client, index) for index in range(_COUNT_BURST)), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Counted)) + + +async def _fire_messages(url: str, key: str) -> tuple[_Sent, ...]: + async def one(client: httpx.AsyncClient) -> _Sent: + marker: Final = uuid.uuid4().hex + body: Final = messages_body(_MODEL, [pdf_document((LETTER,))], f"{_ASK} marker-{marker}") + response: Final = await client.post("/v1/messages", json=body, headers={"Authorization": f"Bearer {key}"}) + return _Sent(marker, response.status_code, response.text, response.headers.get("x-litellm-call-id", "")) + + async with httpx.AsyncClient(base_url=url, timeout=90, trust_env=False) as client: + return tuple(await asyncio.gather(*(one(client) for _ in range(_MESSAGE_BURST)))) + + +def _baseline(gateway: Gateway) -> int: + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _count_body(0)) + assert response.status_code == 200, response.text + counted: Final = _input_tokens(response) + assert counted is not None, response.text + return counted + + +def _assert_exact(counts: tuple[_Counted, ...], baseline: int) -> None: + for item in counts: + assert item.status == 200, (item.pages, item.status, item.text) + assert item.input_tokens is not None and item.input_tokens - baseline == rendered_tokens( + (LETTER,) * item.pages + ), ( + item.pages, + item.input_tokens, + baseline, + ) + + +def _single_spend_row(item: _Sent) -> None: + assert item.status == 200, (item.status, item.text) + assert item.call_id, item.text + rows: Final = eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (item.call_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + assert len(rows) == 1, item.call_id + + +def _held_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 180) +async def test_pdf_count_burst_across_two_workers_prices_every_page_and_logs_each_message_once( + gateway: Gateway, tmp_path: Path +) -> None: + with wire_server(_peer) as wire: + config: Final = owned_config(tmp_path, [_deployment(wire.url)], litellm_settings={"cache": False}) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + eventually( + lambda: _STARTED_WORKER.findall(owned.log.read_text()), + lambda pids: len(pids) == 2, + seconds=graceful_stop_seconds(), + ) + owned_url: Final = str(owned.gateway.client.base_url) + baseline: Final = _baseline(owned.gateway) + wire.drain() + counts, messages = await asyncio.gather( + _fire_counts(owned_url, owned.gateway.key), _fire_messages(owned_url, owned.gateway.key) + ) + received: Final = wire.drain() + targets: Final = [request.target for request in received] + assert targets.count(COUNT_TOKENS_TARGET) == _COUNT_BURST, targets + assert targets.count("/v1/messages") == _MESSAGE_BURST, targets + _assert_exact(counts, baseline) + assert {item.marker for item in messages} == { + _JSON_OBJECT.validate_json(item.text)["id"][len("msg_") :] for item in messages if item.status == 200 + } + for item in messages: + _single_spend_row(item) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 180) +async def test_worker_sigkill_mid_pdf_count_burst_leaves_the_sibling_pricing_pages( + gateway: Gateway, tmp_path: Path +) -> None: + release: Final = threading.Event() + with wire_server(_held(release)) as wire: + config: Final = owned_config(tmp_path, [_deployment(wire.url)], litellm_settings={"cache": False}) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + workers: Final[tuple[int, ...]] = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=graceful_stop_seconds(), + ) + owned_url: Final = str(owned.gateway.client.base_url) + burst: Final = asyncio.create_task( + _fire_counts(owned_url, owned.gateway.key, tolerate_transport_errors=True) + ) + try: + await asyncio.to_thread( + eventually, lambda: wire.received.qsize(), lambda size: size >= _COUNT_BURST, 90 + ) + held_by: Final = MappingProxyType({pid: _held_upstream_connections(pid, wire.url) for pid in workers}) + victim: Final = max(workers, key=held_by.__getitem__) + os.kill(victim, signal.SIGKILL) + finally: + release.set() + served: Final = await burst + wire.drain() + baseline: Final = _baseline(owned.gateway) + after: Final = await _fire_counts(owned_url, owned.gateway.key) + after_received: Final = wire.drain() + assert sum(held_by.values()) == _COUNT_BURST, held_by + assert held_by[victim] > 0, held_by + assert len(served) == _COUNT_BURST - held_by[victim], (len(served), held_by) + assert len(after) == _COUNT_BURST, len(after) + assert [request.target for request in after_received] == [COUNT_TOKENS_TARGET] * (_COUNT_BURST + 1) + _assert_exact(served, baseline) + _assert_exact(after, baseline) diff --git a/tests/integration/providers/test_pdf_document_count_tokens_wire.py b/tests/integration/providers/test_pdf_document_count_tokens_wire.py new file mode 100644 index 00000000000..41d5f3d8685 --- /dev/null +++ b/tests/integration/providers/test_pdf_document_count_tokens_wire.py @@ -0,0 +1,273 @@ +import io +import json +import uuid +from collections.abc import Sequence +from types import MappingProxyType +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.pdf_document import ( + COUNT_REFUSED, + COUNT_TOKENS_TARGET, + LETTER, + NARROW, + Page, + base64_source, + chat_file, + document, + encoded, + messages_body, + pdf_bytes, + pdf_data_url, + pdf_document, + rendered_tokens, + responses_input_file, + text_document, +) +from integration._support.provider import SharedProvider +from integration._support.wire import Reply +from pydantic import JsonValue, TypeAdapter +from pypdf import PdfReader, PdfWriter + +_MODEL: Final = "anthropic/claude-opus-5-5" +_ASK: Final = "Summarize the attached report in one sentence." +_TEXT: Final = "The quarterly report covers revenue, margins and headcount." +_TITLE: Final = "Quarterly report" +_CONTEXT: Final = "Board pack, page one" +_PEER_INPUT_TOKENS: Final = 12 +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_PNG_HEADER: Final = b"\x89PNG\r\n\x1a\n\x00\x00\x00\x0dIHDR" + (1568).to_bytes(4, "big") * 2 + b"\x08\x06\x00\x00\x00" + + +def _locked(raw: bytes, user_password: str) -> bytes: + writer: Final = PdfWriter(clone_from=PdfReader(io.BytesIO(raw))) + writer.encrypt(user_password=user_password, owner_password="owner", algorithm="RC4-128") + out: Final = io.BytesIO() + writer.write(out) + return out.getvalue() + + +_UNREADABLE: Final[MappingProxyType[str, JsonValue]] = MappingProxyType( + { + "garbage-5kb": encoded(bytes(range(256)) * 20), + "cut-mid-stream": encoded(pdf_bytes((LETTER, LETTER))[:200]), + "png-bytes": encoded(_PNG_HEADER), + "user-password": encoded(_locked(pdf_bytes((LETTER, LETTER)), "reader")), + "empty": "", + "int": 1234, + } +) + + +def _int(value: JsonValue) -> int: + assert isinstance(value, int), value + return value + + +def _turn(blocks: Sequence[JsonValue]) -> list[JsonValue]: + return messages_body(_MODEL, blocks, _ASK)["messages"] + + +def _count(gateway: Gateway, provider: SharedProvider, blocks: Sequence[JsonValue]) -> int: + provider.expect(COUNT_REFUSED) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": _MODEL, "messages": _turn(blocks)}) + assert response.status_code == 200, response.text + assert [(request.method, request.target) for request in provider.received()] == [("POST", COUNT_TOKENS_TARGET)] + return _int(_JSON_OBJECT.validate_json(response.content)["input_tokens"]) + + +def _local(gateway: Gateway, parts: Sequence[JsonValue]) -> int: + response: Final = gateway.request( + "POST", "/utils/token_counter", {"model": _MODEL, "messages": _turn(parts)}, params={"call_endpoint": "false"} + ) + assert response.status_code == 200, response.text + return _int(_JSON_OBJECT.validate_json(response.content)["total_tokens"]) + + +def _responses_count(gateway: Gateway, provider: SharedProvider, items: Sequence[JsonValue]) -> int: + provider.expect(COUNT_REFUSED) + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": _MODEL, "input": [{"role": "user", "content": [*items, {"type": "input_text", "text": _ASK}]}]}, + ) + assert response.status_code == 200, response.text + assert [(request.method, request.target) for request in provider.received()] == [("POST", COUNT_TOKENS_TARGET)] + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["object"] == "response.input_tokens", response.text + return _int(payload["input_tokens"]) + + +def _cost(gateway: Gateway, parts: Sequence[JsonValue]) -> float: + payload: Final = gateway.post("/spend/calculate", {"model": _MODEL, "messages": _turn(parts)}) + cost: Final = payload["cost"] + assert isinstance(cost, (int, float)), payload + return float(cost) + + +def _input_price(gateway: Gateway) -> float: + listed: Final = gateway.get("/model/info")["data"] + assert isinstance(listed, list), listed + rows: Final = [object_value(row) for row in listed if isinstance(row, dict) and row.get("model_name") == _MODEL] + assert len(rows) == 1, rows + price: Final = object_value(rows[0]["model_info"])["input_cost_per_token"] + assert isinstance(price, float) and price > 0, price + return price + + +def _peer_message(identity: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": identity, + "type": "message", + "role": "assistant", + "model": "claude-opus-5-5", + "content": [{"type": "text", "text": "One sentence."}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": _PEER_INPUT_TOKENS, "output_tokens": 3}, + } + ).encode() + ) + + +def _prompt_tokens(call_id: str) -> int: + rows: Final = eventually( + lambda: read_rows('SELECT prompt_tokens FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (call_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + return _int(rows[0]["prompt_tokens"]) + + +@pytest.mark.parametrize("pages", [1, 3, 12]) +def test_count_tokens_fallback_prices_each_blank_page_as_a_rendered_image( + gateway: Gateway, provider: SharedProvider, pages: int +) -> None: + letter_pages: Final = (LETTER,) * pages + with_document: Final = _count(gateway, provider, [pdf_document(letter_pages)]) + without: Final = _count(gateway, provider, []) + assert with_document - without == rendered_tokens(letter_pages), (with_document, without) + + +def test_count_tokens_fallback_adds_a_page_text_to_its_rendered_image( + gateway: Gateway, provider: SharedProvider +) -> None: + text_page: Final = _count(gateway, provider, [pdf_document((Page(text=_TEXT),))]) + blank_page: Final = _count(gateway, provider, [pdf_document((LETTER,))]) + text_source: Final = _count(gateway, provider, [text_document(_TEXT)]) + without: Final = _count(gateway, provider, []) + assert text_source - without > 0, (text_source, without) + assert text_page - blank_page == text_source - without, (text_page, blank_page, text_source, without) + + +def test_count_tokens_fallback_prices_a_page_by_its_rendered_area(gateway: Gateway, provider: SharedProvider) -> None: + mixed: Final = (LETTER, NARROW, LETTER) + without: Final = _count(gateway, provider, []) + narrow: Final = _count(gateway, provider, [pdf_document((NARROW,))]) + assert narrow - without == rendered_tokens((NARROW,)), (narrow, without) + assert _count(gateway, provider, [pdf_document(mixed)]) - without == rendered_tokens(mixed) + + +def test_count_tokens_fallback_prices_every_document_in_the_turn(gateway: Gateway, provider: SharedProvider) -> None: + both: Final = _count(gateway, provider, [pdf_document((LETTER,) * 3), pdf_document((LETTER,))]) + without: Final = _count(gateway, provider, []) + assert both - without == rendered_tokens((LETTER,) * 4), (both, without) + + +def test_count_tokens_fallback_prices_a_pdf_without_pages_as_nothing( + gateway: Gateway, provider: SharedProvider +) -> None: + assert _count(gateway, provider, [pdf_document(())]) == _count(gateway, provider, []) + + +def test_count_tokens_fallback_adds_title_and_context_once_per_document( + gateway: Gateway, provider: SharedProvider +) -> None: + without: Final = _count(gateway, provider, []) + plain: Final = _count(gateway, provider, [pdf_document((LETTER,))]) + annotated: Final = _count(gateway, provider, [pdf_document((LETTER,), title=_TITLE, context=_CONTEXT)]) + title: Final = _count(gateway, provider, [text_document(_TITLE)]) - without + context: Final = _count(gateway, provider, [text_document(_CONTEXT)]) - without + assert title > 0 and context > 0, (title, context) + assert annotated - plain == title + context, (annotated, plain, title, context) + + +@pytest.mark.parametrize("label", list(_UNREADABLE)) +def test_count_tokens_fallback_prices_unreadable_pdf_bytes_like_one_image( + gateway: Gateway, provider: SharedProvider, label: str +) -> None: + data: Final = _UNREADABLE[label] + as_pdf: Final = _count(gateway, provider, [document(base64_source(data))]) + as_png: Final = _count(gateway, provider, [document(base64_source(data, media_type="image/png"))]) + without: Final = _count(gateway, provider, []) + assert as_pdf == as_png, (as_pdf, as_png) + assert 0 <= as_pdf - without < rendered_tokens((LETTER,)), (as_pdf, without) + + +def test_count_tokens_fallback_reads_a_list_wrapped_base64_string_like_the_bare_string( + gateway: Gateway, provider: SharedProvider +) -> None: + raw: Final = encoded(pdf_bytes((LETTER,))) + wrapped: Final = _count(gateway, provider, [document(base64_source([raw]))]) + assert wrapped == _count(gateway, provider, [document(base64_source(raw))]), wrapped + + +def test_count_tokens_fallback_reads_an_owner_locked_pdf(gateway: Gateway, provider: SharedProvider) -> None: + pages: Final = (LETTER, LETTER) + locked: Final = document(base64_source(encoded(_locked(pdf_bytes(pages), "")))) + assert _count(gateway, provider, [locked]) - _count(gateway, provider, []) == rendered_tokens(pages) + + +def test_count_tokens_fallback_parses_only_the_pdf_media_type(gateway: Gateway, provider: SharedProvider) -> None: + raw: Final = encoded(pdf_bytes((LETTER,))) + without: Final = _count(gateway, provider, []) + labelled: Final = _count(gateway, provider, [document(base64_source(raw))]) + assert labelled - without == rendered_tokens((LETTER,)), (labelled, without) + mislabelled: Final = _count(gateway, provider, [document(base64_source(raw, media_type="application/x-pdf"))]) + assert mislabelled == _count(gateway, provider, [document(base64_source(raw, media_type="image/png"))]) + + +def test_utils_token_counter_prices_a_chat_file_by_its_pages(gateway: Gateway) -> None: + one: Final = _local(gateway, [chat_file(pdf_data_url((LETTER,)))]) + three: Final = _local(gateway, [chat_file(pdf_data_url((LETTER,) * 3))]) + assert three - one == rendered_tokens((LETTER,) * 2), (three, one) + text_page: Final = _local(gateway, [chat_file(pdf_data_url((Page(text=_TEXT),)))]) + text_part: Final = _local(gateway, [chat_file(pdf_data_url((LETTER,))), {"type": "text", "text": _TEXT}]) + assert text_part - one > 0, (text_part, one) + assert text_page - one == text_part - one, (text_page, text_part, one) + + +def test_responses_input_tokens_fallback_prices_an_input_file_by_its_pages( + gateway: Gateway, provider: SharedProvider +) -> None: + one: Final = _responses_count(gateway, provider, [responses_input_file(pdf_data_url((LETTER,)))]) + three: Final = _responses_count(gateway, provider, [responses_input_file(pdf_data_url((LETTER,) * 3))]) + assert three - one == rendered_tokens((LETTER,) * 2), (three, one) + + +def test_spend_calculate_prices_a_chat_file_by_its_pages(gateway: Gateway) -> None: + twelve: Final = _cost(gateway, [chat_file(pdf_data_url((LETTER,) * 12))]) + one: Final = _cost(gateway, [chat_file(pdf_data_url((LETTER,)))]) + assert twelve - one == pytest.approx(rendered_tokens((LETTER,) * 11) * _input_price(gateway)), (twelve, one) + + +def test_a_cached_pdf_message_keeps_the_peer_usage_on_both_spend_rows( + gateway: Gateway, provider: SharedProvider +) -> None: + marker: Final = uuid.uuid4().hex + body: Final = messages_body(_MODEL, [pdf_document((LETTER,))], f"{_ASK} marker-{marker}") + provider.expect(_peer_message(f"msg_{marker}")) + first: Final = gateway.request("POST", "/v1/messages", body) + second: Final = gateway.request("POST", "/v1/messages", body) + assert (first.status_code, second.status_code) == (200, 200), (first.text, second.text) + assert [(request.method, request.target) for request in provider.received()] == [("POST", "/v1/messages")] + assert _JSON_OBJECT.validate_json(second.content)["id"] == f"msg_{marker}", second.text + assert first.headers["x-litellm-call-id"] != second.headers["x-litellm-call-id"] + assert [_prompt_tokens(response.headers["x-litellm-call-id"]) for response in (first, second)] == [ + _PEER_INPUT_TOKENS, + _PEER_INPUT_TOKENS, + ] diff --git a/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py b/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py new file mode 100644 index 00000000000..9fad73f8080 --- /dev/null +++ b/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py @@ -0,0 +1,292 @@ +import os +import socket +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment, object_value +from integration._support.database import read_rows, scratch_database +from integration._support.pdf_document import ( + COUNT_REFUSED, + COUNT_TOKENS_TARGET, + LETTER, + chat_body, + chat_file, + messages_body, + pdf_data_url, + pdf_document, + responses_body, + responses_input_file, +) +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import anthropic_peer, marker_of, owned_config +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +_WINDOW: Final = "window-claude" +_ITPM: Final = "itpm-claude" +_BUDGET: Final = "budget-claude" +_AFFINITY: Final = "affinity-claude" +_LIMIT: Final = 5000 +_PAGE_COUNT: Final = 12 +_PAGES: Final = (LETTER,) * _PAGE_COUNT +_ASK: Final = "Summarize the attached report in one sentence." +_FOLLOW_UPS: Final = 9 +_TRANSCRIPT: Final = " ".join(f"Line {number}: revenue, margins and headcount moved." for number in range(160)) +_PIN_KEYS: Final = "*:prompt_caching" +_PROVIDER_KEY: Final = "synthetic-provider-key" +_SURFACES: Final = ("messages", "messages-stream", "chat-file", "responses-input-file") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + wire: Wire + held_port: int + + +def _marker() -> str: + return uuid.uuid4().hex + + +def _free_port() -> int: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + return int(reserve.getsockname()[1]) + + +def _deployment( + name: str, + backend: str, + api_base: str, + *, + model_info: Mapping[str, JsonValue] | None = None, + **params: JsonValue, +) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": f"anthropic/{backend}", "api_base": api_base, "api_key": _PROVIDER_KEY, **params}, + "model_info": dict(model_info or {}), + } + + +def _peer(request: Request) -> Reply: + if request.target == COUNT_TOKENS_TARGET: + return COUNT_REFUSED + return anthropic_peer(request) + + +def _held(release: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert release.wait(timeout=120), "Held peer was never released" + return _peer(request) + + return respond + + +def _request(surface: str, model: str, pages: int, marker: str) -> tuple[str, dict[str, JsonValue]]: + text: Final = f"{_ASK} marker-{marker}" + letter_pages: Final = (LETTER,) * pages + if surface == "messages": + return "/v1/messages", messages_body(model, [pdf_document(letter_pages)], text) + if surface == "messages-stream": + return "/v1/messages", messages_body(model, [pdf_document(letter_pages)], text, stream=True) + if surface == "chat-file": + return "/v1/chat/completions", chat_body(model, [chat_file(pdf_data_url(letter_pages))], text) + return "/v1/responses", responses_body(model, [responses_input_file(pdf_data_url(letter_pages))], text) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("pdf-pre-call") + held_port: Final = _free_port() + with gateway_from_environment() as gateway, wire_server(_peer) as wire: + config: Final = owned_config( + directory, + [ + _deployment(_WINDOW, "claude-haiku-5-5", wire.url, model_info={"max_input_tokens": _LIMIT}), + _deployment(_ITPM, "claude-sonnet-4-6", wire.url, itpm=_LIMIT), + _deployment(_BUDGET, "claude-opus-4-8", f"http://127.0.0.1:{held_port}"), + ], + litellm_settings={"cache": False}, + router_settings={ + "enable_pre_call_checks": True, + "optional_pre_call_checks": ["enforce_model_rate_limits"], + }, + ) + with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned: + yield _Rig(owned, wire, held_port) + + +def _spend_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows('SELECT status FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (call_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + return rows[0] + + +@pytest.mark.parametrize("surface", _SURFACES) +def test_context_window_check_rejects_a_pdf_whose_pages_exceed_max_input_tokens(rig: _Rig, surface: str) -> None: + rig.wire.drain() + path, body = _request(surface, _WINDOW, _PAGE_COUNT, _marker()) + response: Final = rig.owned.gateway.request("POST", path, body) + assert response.status_code == 400, (response.status_code, response.text) + assert "Context Window exceeded" in response.text, response.text + assert [request.target for request in rig.wire.drain()] == [] + + +def test_context_window_check_passes_a_one_page_pdf_and_forwards_its_bytes(rig: _Rig) -> None: + rig.wire.drain() + path, body = _request("messages", _WINDOW, 1, _marker()) + response: Final = rig.owned.gateway.request("POST", path, body) + assert response.status_code == 200, response.text + received: Final = rig.wire.drain() + assert [request.target for request in received] == ["/v1/messages"] + forwarded: Final = _JSON_OBJECT.validate_json(received[0].body)["messages"] + assert isinstance(forwarded, list) and len(forwarded) == 1, forwarded + blocks: Final = object_value(forwarded[0])["content"] + assert isinstance(blocks, list) and pdf_document((LETTER,)) in blocks, blocks + + +def test_model_itpm_check_rejects_a_pdf_whose_pages_exceed_the_limit(rig: _Rig) -> None: + rig.wire.drain() + path, body = _request("messages", _ITPM, _PAGE_COUNT, _marker()) + response: Final = rig.owned.gateway.request("POST", path, body) + assert response.status_code == 429, (response.status_code, response.text) + assert '"error"' in response.text, response.text + assert [request.target for request in rig.wire.drain()] == [] + + +def test_key_budget_reservation_rejects_the_second_concurrent_pdf_request(rig: _Rig) -> None: + release: Final = threading.Event() + requests: Final = tuple(_request("messages", _BUDGET, _PAGE_COUNT, _marker()) for _ in range(2)) + with wire_server(_held(release), port=rig.held_port) as wire, rig.owned.gateway.scenario() as scenario: + key: Final = scenario.key(max_budget=0.01) + headers: Final = {"Authorization": f"Bearer {key}"} + with ( + httpx.Client(base_url=str(rig.owned.gateway.client.base_url), timeout=90, trust_env=False) as client, + ThreadPoolExecutor(max_workers=2) as pool, + ): + futures: Final = tuple( + pool.submit(client.post, path, json=body, headers=headers) for path, body in requests + ) + try: + eventually( + lambda: wire.received.qsize() + sum(1 for future in futures if future.done()), + lambda settled: settled >= 2, + seconds=60, + ) + finally: + release.set() + responses: Final = tuple(future.result(timeout=90) for future in futures) + received: Final = wire.drain() + assert sorted(response.status_code for response in responses) == [200, 422], [ + response.text for response in responses + ] + rejected: Final = next(response for response in responses if response.status_code != 200) + assert "udget" in rejected.text, rejected.text + assert [request.target for request in received] == ["/v1/messages"] + + +def test_a_pre_call_rejection_logs_one_failure_row_and_calls_no_upstream(rig: _Rig) -> None: + path, body = _request("messages", _WINDOW, _PAGE_COUNT, _marker()) + response: Final = rig.owned.gateway.request("POST", path, body) + assert response.status_code == 400, (response.status_code, response.text) + assert _spend_row(response.headers["x-litellm-call-id"])["status"] == "failure" + assert rig.wire.drain() == () + + +def _first_turn(session: str, marker: str) -> list[JsonValue]: + return [ + {"role": "system", "content": f"Answer from the attached report. session {session}"}, + *chat_body( + _AFFINITY, + [{"type": "text", "text": _TRANSCRIPT}, chat_file(pdf_data_url(_PAGES))], + f"{_ASK} marker-{marker}", + )["messages"], + ] + + +def _follow_up(first_turn: list[JsonValue], marker: str) -> list[JsonValue]: + return [ + *first_turn, + {"role": "assistant", "content": "One sentence."}, + {"role": "user", "content": f"And the margins? marker-{marker}"}, + ] + + +def _served(wire: Wire) -> frozenset[str]: + return frozenset(marker_of(request) for request in wire.drain()) + + +def _pin_count(cache: Redis) -> int: + return len(cache.keys(_PIN_KEYS)) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 180) +def test_prompt_caching_affinity_pins_follow_ups_of_a_cached_turn_that_carries_a_pdf( + gateway: Gateway, tmp_path: Path +) -> None: + session: Final = _marker() + first_marker: Final = _marker() + follow_up_markers: Final = tuple(_marker() for _ in range(_FOLLOW_UPS)) + first_turn: Final = _first_turn(session, first_marker) + with scratch_database() as database_url, wire_server(_peer) as left, wire_server(_peer) as right: + config: Final = owned_config( + tmp_path, + [ + _deployment(_AFFINITY, "claude-opus-5-5", left.url, model_info={"id": f"pdf-left-{session}"}), + _deployment(_AFFINITY, "claude-opus-5-5", right.url, model_info={"id": f"pdf-right-{session}"}), + ], + litellm_settings={"cache": False}, + router_settings={ + "optional_pre_call_checks": ["prompt_caching"], + "redis_host": os.environ["REDIS_HOST"], + "redis_port": int(os.environ["REDIS_PORT"]), + }, + ) + with ( + Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache, + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": database_url}, + config=config, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=2, + ) as owned, + owned.gateway.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"enable_prompt_caching": True}) + pins_before: Final = _pin_count(cache) + first: Final = owned.gateway.request( + "POST", "/v1/chat/completions", {"model": _AFFINITY, "messages": first_turn, "max_tokens": 64}, key=key + ) + assert first.status_code == 200, first.text + eventually(lambda: _pin_count(cache), lambda count: count > pins_before, seconds=70) + follow_ups: Final = tuple( + owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": _AFFINITY, "messages": _follow_up(first_turn, marker), "max_tokens": 64}, + key=key, + ) + for marker in follow_up_markers + ) + served: Final = {"left": _served(left), "right": _served(right)} + assert [response.status_code for response in follow_ups] == [200] * _FOLLOW_UPS, [ + response.text for response in follow_ups + ] + assert served["left"] | served["right"] == {first_marker, *follow_up_markers}, served + first_side: Final = next(side for side in served if first_marker in served[side]) + assert served[first_side] == {first_marker, *follow_up_markers}, served diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index c71b1496bdd..42c25a319a2 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -11,6 +11,7 @@ import threading import time from collections.abc import Mapping from concurrent.futures import Future, wait +from itertools import accumulate, chain from pathlib import Path from typing import Final from unittest.mock import MagicMock @@ -1444,7 +1445,7 @@ def _count_user_content(content: list[dict]) -> int: ids=["base64", "url", "file"], ) def test_anthropic_document_block_with_opaque_source_is_priced_like_an_image(source: dict[str, str]): - """A `document` whose bytes can't be tokenized locally is priced like an `image`, not raised on.""" + """A `document` with no readable pages (a bare PDF header, a URL, a file id) is priced like an `image`, not raised on.""" prompt = {"type": "text", "text": "Summarize this file."} assert _count_user_content([prompt, {"type": "document", "source": source}]) == _count_user_content( @@ -1504,6 +1505,12 @@ def test_openai_file_block_prices_like_the_equivalent_anthropic_document(): assert _count_user_content([prompt, inline_file]) == _count_user_content([prompt, document]) assert _count_user_content([prompt, inline_file]) > _count_user_content([prompt]) + readable: Final = _pdf_base64(("Revenue grew eleven percent while churn fell to two percent.",)) + readable_file: Final = {"type": "file", "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64," + readable}} + readable_document: Final = {"type": "document", "title": "report.pdf", "source": _pdf_source(readable)} + assert _count_user_content([prompt, readable_file]) == _count_user_content([prompt, readable_document]) + assert _count_user_content([prompt, readable_file]) > _count_user_content([prompt, inline_file]) + def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): """A `file` block naming an uploaded file has no bytes to price, so it adds only the filename's tokens.""" @@ -1518,6 +1525,101 @@ def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): ) +def _pdf_base64(pages: tuple[str, ...], width: int = 612, height: int = 792) -> str: + def page_objects(index: int, text: str) -> tuple[bytes, bytes]: + escaped: Final = text.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)") + stream: Final = f"BT /F1 12 Tf 72 720 Td ({escaped}) Tj ET".encode("latin-1") + content: Final = f"<< /Length {len(stream)} >>\nstream\n".encode() + stream + b"\nendstream" + page: Final = ( + f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {width} {height}] " + f"/Resources << /Font << /F1 3 0 R >> >> /Contents {4 + 2 * index} 0 R >>" + ).encode() + return content, page + + kids: Final = " ".join(f"{5 + 2 * index} 0 R" for index in range(len(pages))) + bodies: Final = ( + b"<< /Type /Catalog /Pages 2 0 R >>", + f"<< /Type /Pages /Count {len(pages)} /Kids [ {kids} ] >>".encode(), + b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>", + *chain.from_iterable(page_objects(index, text) for index, text in enumerate(pages)), + ) + header: Final = b"%PDF-1.4\n" + objects: Final = tuple( + f"{number} 0 obj\n".encode() + body + b"\nendobj\n" for number, body in enumerate(bodies, start=1) + ) + offsets: Final = accumulate((len(header), *(len(obj) for obj in objects[:-1]))) + xref: Final = f"xref\n0 {len(objects) + 1}\n0000000000 65535 f \n".encode() + b"".join( + f"{offset:010d} 00000 n \n".encode() for offset in offsets + ) + trailer: Final = ( + f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\n" + f"startxref\n{len(header) + sum(len(obj) for obj in objects)}\n%%EOF\n" + ).encode() + return base64.b64encode(header + b"".join(objects) + xref + trailer).decode() + + +def _pdf_source(pdf_base64: str) -> dict[str, str]: + return {"type": "base64", "media_type": "application/pdf", "data": pdf_base64} + + +def test_base64_pdf_document_counts_every_page_text_and_rendering(): + """A base64 PDF is read page by page: each page costs its text plus the image Anthropic renders it to. + + Before the fix the whole document was priced as one 85-token image, so a count_tokens call that fell back + to the local counter answered 116 for a 12-page PDF the provider then billed at 35941 input tokens. + """ + prompt: Final = {"type": "text", "text": "Summarize this file."} + first: Final = "Revenue grew eleven percent while churn fell to two percent." + second: Final = "Headcount is flat and the office lease was renewed for three years." + base: Final = _count_user_content([prompt]) + blank_page: Final = _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64(("",)))}]) - base + + assert blank_page > 0 + assert _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64((first, second)))}]) == ( + _count_user_content([prompt, {"type": "text", "text": first}, {"type": "text", "text": second}]) + 2 * blank_page + ) + one_page: Final = _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64((first,)))}]) - base + twelve_pages: Final = ( + _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64((first,) * 12))}]) - base + ) + assert one_page > _count_user_content([prompt, {"type": "text", "text": first}]) - base + assert twelve_pages == 12 * one_page + + +@pytest.mark.parametrize("fields", [{}, {"title": "Q3 board packet", "context": "Shared by finance"}], ids=["bare", "described"]) +def test_pdf_page_rendering_cost_follows_anthropic_image_scaling(fields: dict[str, str]): + """Anthropic rasterizes each PDF page within its image limits (1568 px long edge, 1.15 MP) and bills + width * height / 750 tokens for it: https://platform.claude.com/docs/en/build-with-claude/pdf-support and + https://platform.claude.com/docs/en/build-with-claude/vision, read 2026-10-07, when Bedrock billed about + 1550 tokens per blank Letter page on both Sonnet 4.6 and Opus 4.8. + """ + prompt: Final = {"type": "text", "text": "Summarize this file."} + + def blank_page_cost(width: int, height: int) -> int: + with_page: Final = {"type": "document", "source": _pdf_source(_pdf_base64(("",), width, height)), **fields} + without_page: Final = {"type": "document", "source": _pdf_source(_pdf_base64((), width, height)), **fields} + return _count_user_content([prompt, with_page]) - _count_user_content([prompt, without_page]) + + letter: Final = blank_page_cost(612, 792) + poster: Final = blank_page_cost(2448, 3168) + strip: Final = blank_page_cost(1000, 100) + + assert letter == poster == 1534 + assert strip == 328 + + +def test_pdf_document_without_pypdf_is_priced_like_an_image(monkeypatch: pytest.MonkeyPatch): + prompt: Final = {"type": "text", "text": "Summarize this file."} + source: Final = _pdf_source(_pdf_base64(("Revenue grew eleven percent while churn fell to two percent.",))) + priced_by_page: Final = _count_user_content([prompt, {"type": "document", "source": source}]) + + monkeypatch.setitem(sys.modules, "pypdf", None) + + priced_as_image: Final = _count_user_content([prompt, {"type": "document", "source": source}]) + assert priced_as_image == _count_user_content([prompt, {"type": "image", "source": source}]) + assert priced_as_image < priced_by_page + + def _png_data_url(width: int, height: int) -> str: ihdr = b"\x89PNG\r\n\x1a\n" + (13).to_bytes(4, "big") + b"IHDR" + width.to_bytes(4, "big") + height.to_bytes(4, "big") return "data:image/png;base64," + base64.b64encode(ihdr + b"\x08\x06\x00\x00\x00").decode() From 095c4eaa41e5e1f7de530503e49c12648fdb3019 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:14:05 -0700 Subject: [PATCH 10/36] fix(anthropic): count a leading system run through count_tokens' system parameter (#45463) Backport of #45463 to stable/1.104.x. Cherry-picked from d34ad3528191e2d338b6f8f982cf97482ef52bf3 (main). Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../mid_conversation_system.py | 16 +- .../anthropic/count_tokens/transformation.py | 40 +- .../providers/_count_tokens_system_lift.py | 159 +++++ ...anthropic_count_tokens_system_lift_wire.py | 572 ++++++++++++++++++ .../test_bedrock_mantle_count_tokens_wire.py | 368 ++++++++++- ...rompt_templates_mid_conversation_system.py | 21 + ...t_anthropic_count_tokens_transformation.py | 112 ++++ 7 files changed, 1254 insertions(+), 34 deletions(-) create mode 100644 tests/integration/providers/_count_tokens_system_lift.py create mode 100644 tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py index b5e9afca86b..37c32144bbf 100644 --- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py +++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py @@ -31,7 +31,7 @@ Anthropic wire shape is built later by ``anthropic_messages_pt``. from collections.abc import Iterator, Mapping, Sequence from itertools import chain, groupby -from typing import Final, Literal, TypeAlias +from typing import Final, Literal, TypeAlias, TypeVar from litellm.types.llms.anthropic import AnthropicMessagesSystemMessageParam, AnthropicSystemMessageContent from litellm.types.llms.openai import ( @@ -55,6 +55,7 @@ _RENDERED_ASSISTANT_PART_TYPES: Final = frozenset({"text", "server_tool_use"}) _THINKING_BLOCK_TYPES: Final = frozenset({"thinking", "redacted_thinking"}) _MessageKind: TypeAlias = Literal["system", "tool", "user", "other"] +_Message: Final = TypeVar("_Message") _TextPart: TypeAlias = tuple[str, ChatCompletionCachedContent | None] @@ -96,8 +97,8 @@ def _kind(message: object) -> _MessageKind: def split_leading_system_run( - messages: Sequence[AllMessageValues], -) -> tuple[tuple[AllMessageValues, ...], tuple[AllMessageValues, ...]]: + messages: Sequence[_Message], +) -> tuple[tuple[_Message, ...], tuple[_Message, ...]]: """Split ``messages`` into the leading run of system messages and everything after it.""" leading_count: Final = next( (index for index, message in enumerate(messages) if not is_system_message(message)), @@ -160,9 +161,16 @@ def _anthropic_text_block(part: _TextPart) -> AnthropicSystemMessageContent: return cached +def anthropic_system_blocks(run: Sequence[object]) -> tuple[AnthropicSystemMessageContent, ...]: + """The top-level ``system`` blocks for a run of system messages: every non-empty text part in order, + each keeping its ``cache_control``, which is the shape the chat path sends for the leading run.""" + parts: Final = chain.from_iterable(_text_parts(message) for message in run) + return tuple(_anthropic_text_block(part) for part in parts) + + def anthropic_system_messages(message: object) -> tuple[AnthropicMessagesSystemMessageParam, ...]: """The Anthropic wire message for a system message, or nothing when it carries no text.""" - blocks: Final = tuple(_anthropic_text_block(part) for part in _text_parts(message)) + blocks: Final = anthropic_system_blocks((message,)) if not blocks: return () wire: Final[AnthropicMessagesSystemMessageParam] = { diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 8147496b402..8c47d1cf884 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -11,9 +11,14 @@ from typing import Final from pydantic import JsonValue, TypeAdapter from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION +from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( + anthropic_system_blocks, + split_leading_system_run, +) from litellm.types.llms.openai import ChatCompletionImageObject _COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue]) +_SYSTEM_BLOCKS: Final = TypeAdapter(list[JsonValue]) _IMAGE_BLOCK: Final = TypeAdapter(ChatCompletionImageObject) COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config") @@ -46,6 +51,27 @@ def _count_content(content: JsonValue) -> JsonValue: return [_count_block(block) for block in content] if isinstance(content, list) else content +def _lift_leading_system( + messages: Sequence[Mapping[str, JsonValue]], system: JsonValue +) -> tuple[tuple[Mapping[str, JsonValue], ...], JsonValue]: + """Move the leading run of system-role messages into the top-level ``system`` parameter. + + count_tokens only takes the initial system prompt there and answers 400 on ``role: "system"`` + at the head of ``messages``; the chat path sends the same run as ``system``. A caller's own + ``system`` keeps its place ahead of the lifted blocks, and a ``system`` that is neither text + nor a block list is left as sent, messages included, for the provider to judge. + """ + leading, conversation = split_leading_system_run(messages) + if not leading or not (system is None or isinstance(system, (str, list))): + return tuple(messages), system + lifted: Final = _SYSTEM_BLOCKS.validate_python(list(anthropic_system_blocks(leading))) + if isinstance(system, list): + return conversation, [*system, *lifted] + if isinstance(system, str) and system: + return conversation, [{"type": "text", "text": system}, *lifted] + return conversation, lifted or system + + class AnthropicCountTokensConfig: """ Configuration and transformation logic for Anthropic CountTokens API. @@ -76,16 +102,24 @@ class AnthropicCountTokensConfig: """ Transform request to Anthropic CountTokens format. - Includes optional system and tools fields for accurate token counting. + Includes optional system and tools fields for accurate token counting; a leading run of + system-role messages is counted through ``system``, the only place count_tokens accepts it. """ options: Final[Mapping[str, JsonValue]] = optional_params or MappingProxyType({}) + counted_messages, counted_system = _lift_leading_system(messages, system) return _COUNT_REQUEST.validate_python( MappingProxyType( { "model": model, - "messages": [{**message, "content": _count_content(message["content"])} for message in messages], + "messages": [ + {**message, "content": _count_content(message["content"])} for message in counted_messages + ], **MappingProxyType( - {key: value for key, value in (("system", system), ("tools", tools)) if value is not None} + { + key: value + for key, value in (("system", counted_system), ("tools", tools)) + if value is not None + } ), **MappingProxyType( {key: value for key, value in options.items() if key in COUNT_TOKEN_OPTION_NAMES} diff --git a/tests/integration/providers/_count_tokens_system_lift.py b/tests/integration/providers/_count_tokens_system_lift.py new file mode 100644 index 00000000000..56619d1c963 --- /dev/null +++ b/tests/integration/providers/_count_tokens_system_lift.py @@ -0,0 +1,159 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import openai +from integration._support.client import Gateway +from pydantic import JsonValue, TypeAdapter + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +TOKEN_COUNTING_BETA: Final = "token-counting-2024-11-01" +ANTHROPIC_VERSION: Final = "2023-06-01" +REJECTION: Final = 'messages.0: Unexpected role "system". The Messages API accepts a top-level `system` parameter' + +USER_TEXT: Final = "Count this message" +USER: Final[dict[str, JsonValue]] = {"role": "user", "content": USER_TEXT} +ASSISTANT: Final[dict[str, JsonValue]] = {"role": "assistant", "content": "One."} +FOLLOW_UP: Final[dict[str, JsonValue]] = {"role": "user", "content": "Again"} +INSTRUCTION: Final = "You are a terse assistant" +REMINDER: Final = "Answer in one sentence" +LEADING: Final[dict[str, JsonValue]] = {"role": "system", "content": INSTRUCTION} +LIFTED: Final[dict[str, JsonValue]] = {"type": "text", "text": INSTRUCTION} +MID_SYSTEM: Final[dict[str, JsonValue]] = {"role": "system", "content": REMINDER} +EPHEMERAL: Final[dict[str, JsonValue]] = {"type": "ephemeral"} +ONE_HOUR: Final[dict[str, JsonValue]] = {"type": "ephemeral", "ttl": "1h"} +IMAGE_PART: Final[dict[str, JsonValue]] = { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" + }, +} +FIVE_KB: Final = "Answer in one sentence. " * 214 +CALLER_SYSTEM: Final = "Prefer metric units" +CALLER_BLOCKS: Final[list[JsonValue]] = [ + {"type": "text", "text": CALLER_SYSTEM}, + {"type": "text", "text": "Never guess", "cache_control": EPHEMERAL}, +] +TOOLS: Final[list[JsonValue]] = [ + { + "name": "get_weather", + "description": "Look up the current weather for a city", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string", "description": "City to look up"}}, + "required": ["city"], + }, + } +] + + +@dataclass(frozen=True, slots=True) +class LiftCase: + messages: tuple[dict[str, JsonValue], ...] + system: JsonValue | None + lifted: list[JsonValue] | None + + +LIFT_CASES: Final[Mapping[str, LiftCase]] = MappingProxyType( + { + "string": LiftCase((LEADING, USER), None, [LIFTED]), + "run_with_cache_control": LiftCase( + ( + {"role": "system", "content": INSTRUCTION, "cache_control": ONE_HOUR}, + { + "role": "system", + "content": [ + {"type": "text", "text": REMINDER, "cache_control": EPHEMERAL}, + {"type": "text", "text": ""}, + IMAGE_PART, + ], + }, + USER, + ), + None, + [ + {"type": "text", "text": INSTRUCTION, "cache_control": ONE_HOUR}, + {"type": "text", "text": REMINDER, "cache_control": EPHEMERAL}, + ], + ), + "caller_system_string_first": LiftCase( + (LEADING, USER), CALLER_SYSTEM, [{"type": "text", "text": CALLER_SYSTEM}, LIFTED] + ), + "caller_system_blocks_first": LiftCase((LEADING, USER), CALLER_BLOCKS, [*CALLER_BLOCKS, LIFTED]), + "caller_empty_system_dropped": LiftCase((LEADING, USER), "", [LIFTED]), + "empty_content_dropped": LiftCase(({"role": "system", "content": ""}, USER), None, None), + "image_only_content_dropped": LiftCase(({"role": "system", "content": [IMAGE_PART]}, USER), None, None), + "integer_content_dropped": LiftCase(({"role": "system", "content": 5}, USER), None, None), + "integer_text_part_dropped": LiftCase( + ({"role": "system", "content": [{"type": "text", "text": 7}]}, USER), None, None + ), + "five_kb": LiftCase(({"role": "system", "content": FIVE_KB}, USER), None, [{"type": "text", "text": FIVE_KB}]), + } +) +STRING_CASE: Final = LIFT_CASES["string"] + + +def count_request(model: str, case: LiftCase) -> dict[str, JsonValue]: + return {"model": model, "messages": list(case.messages), **({} if case.system is None else {"system": case.system})} + + +def expected_count_body(model: str, case: LiftCase) -> dict[str, JsonValue]: + conversation: Final[list[JsonValue]] = [message for message in case.messages if message["role"] != "system"] + return {"model": model, "messages": conversation, **({} if case.lifted is None else {"system": case.lifted})} + + +def _opening_message(message: JsonValue) -> bool: + return isinstance(message, dict) and message.get("role") in ("user", "assistant") + + +def _conversation_message(message: JsonValue) -> bool: + return isinstance(message, dict) and message.get("role") in ("user", "assistant", "system") + + +def _anthropic_tool(tool: JsonValue) -> bool: + return isinstance(tool, dict) and isinstance(tool.get("name"), str) and isinstance(tool.get("input_schema"), dict) + + +def accepts_count_body(body: Mapping[str, JsonValue]) -> bool: + # Anthropic, Azure AI Foundry and Bedrock Mantle count_tokens verdicts observed live on 2026-10-09: an empty + # messages list and a system role at messages[0] answer 400 invalid_request_error, a later system role counts + messages: Final = body.get("messages") + tools: Final = body.get("tools", []) + return ( + isinstance(messages, list) + and len(messages) > 0 + and _opening_message(messages[0]) + and all(map(_conversation_message, messages)) + and isinstance(body.get("system", ""), (str, list)) + and isinstance(tools, list) + and all(map(_anthropic_tool, tools)) + ) + + +def anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False), + ) + + +def async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False), + ) 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 new file mode 100644 index 00000000000..13091ccd14a --- /dev/null +++ b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py @@ -0,0 +1,572 @@ +import asyncio +import json +import socket +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import AbstractContextManager, ExitStack +from dataclasses import dataclass +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, cast + +import httpcore +import httpx +import psutil +import pytest +from anthropic.types import MessageParam +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._count_tokens_system_lift import ( + ANTHROPIC_VERSION, + ASSISTANT, + EPHEMERAL, + FOLLOW_UP, + IMAGE_PART, + INSTRUCTION, + JSON_OBJECT, + LEADING, + LIFT_CASES, + LIFTED, + MID_SYSTEM, + REJECTION, + REMINDER, + STRING_CASE, + TOKEN_COUNTING_BETA, + TOOLS, + USER, + USER_TEXT, + LiftCase, + accepts_count_body, + anthropic_client, + async_anthropic_client, + async_openai_client, + count_request, + expected_count_body, + openai_client, +) +from pydantic import JsonValue, TypeAdapter + +_MODEL: Final = "claude-opus-5-5" +_API_KEY: Final = "synthetic-count-tokens-key" +_COUNT: Final = 3131 +_PROVIDER_TOKENIZERS: Final = frozenset({"anthropic_api", "azure_ai_anthropic_api"}) +_SDK_MESSAGES: Final = cast( + list[MessageParam], [LEADING, USER] +) # cast-ok: the SDK types reject the role the proxy lifts +_STRING_BODY: Final = expected_count_body(_MODEL, STRING_CASE) +_PROXY_MODULE: Final = "integration._support.proxy" +_PROBES_PER_ROUND: Final = 8 +_CLIENT_ADDRESS: Final = TypeAdapter(tuple[str, int]) +_ARGUMENTS: Final = TypeAdapter(tuple[str, ...]) +_NAME: Final = TypeAdapter(str) + + +@dataclass(frozen=True, slots=True) +class _Provider: + prefix: str + target: str + tokenizer: str + azure: bool + + +@dataclass(frozen=True, slots=True) +class _Deployment: + provider: _Provider + port: int + model: str + + +_ANTHROPIC: Final = _Provider("anthropic", "/v1/messages/count_tokens", "anthropic_api", False) +_AZURE: Final = _Provider("azure_ai", "/anthropic/v1/messages/count_tokens", "azure_ai_anthropic_api", True) +_PROVIDERS: Final = MappingProxyType({"anthropic": _ANTHROPIC, "azure_ai": _AZURE}) + + +@pytest.fixture(params=_PROVIDERS.keys()) +def provider(request: pytest.FixtureRequest) -> _Provider: + name: Final[object] = request.param # pyright: ignore[reportAny] # pytest types the fixture param as Any + return _PROVIDERS[_NAME.validate_python(name)] + + +def _json_reply(status: int, payload: Mapping[str, JsonValue]) -> Reply: + return Reply(status=status, body=json.dumps(payload).encode()) + + +def _rejected(status: int) -> Reply: + return _json_reply(status, {"type": "error", "error": {"type": "invalid_request_error", "message": REJECTION}}) + + +def _counted(request: Request) -> Reply: + accepted: Final = accepts_count_body(JSON_OBJECT.validate_json(request.body)) + return _json_reply(200, {"input_tokens": _COUNT}) if accepted else _rejected(400) + + +def _rejecting(status: int) -> Callable[[Request], Reply]: + def count(_request: Request) -> Reply: + return _rejected(status) + + return count + + +def _holding(held: SimpleQueue[str], release: threading.Event, seconds: float) -> Callable[[Request], Reply]: + def hold(request: Request) -> Reply: + held.put(request.target) + assert release.wait(timeout=seconds), "Held count was never released" + return _counted(request) + + return hold + + +def _peer(provider: _Provider, count: Callable[[Request], Reply] = _counted) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == provider.target: + return count(request) + return _json_reply(404, {"error": f"unscripted target {request.target}"}) + + return respond + + +def _listening(deployment: _Deployment, count: Callable[[Request], Reply] = _counted) -> AbstractContextManager[Wire]: + return wire_server(_peer(deployment.provider, count), port=deployment.port) + + +def _reserved_port() -> int: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + return _CLIENT_ADDRESS.validate_python(reserve.getsockname())[1] + + +def _cmdline(process: psutil.Process) -> tuple[str, ...]: + try: + return _ARGUMENTS.validate_python(process.cmdline()) + except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): + return () + + +def _serves(cmdline: tuple[str, ...], proxy_port: int) -> bool: + if _PROXY_MODULE not in cmdline or "--port" not in cmdline: + return False + return cmdline[cmdline.index("--port") + 1] == str(proxy_port) + + +def _listens(process: psutil.Process, proxy_port: int) -> bool: + return any( + connection.status == psutil.CONN_LISTEN and connection.laddr and connection.laddr.port == proxy_port + for connection in process.net_connections(kind="tcp") + ) + + +def _proxy_workers(proxy_port: int) -> frozenset[int]: + (master,) = tuple(process for process in psutil.process_iter() if _serves(_cmdline(process), proxy_port)) + spawned: Final = frozenset(child.pid for child in master.children() if _listens(child, proxy_port)) + return spawned or frozenset({master.pid}) + + +def _holder(workers: frozenset[int], client_port: int) -> int | None: + def holds(pid: int) -> bool: + return any( + connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == client_port + for connection in psutil.Process(pid).net_connections(kind="tcp") + ) + + return next((pid for pid in sorted(workers) if holds(pid)), None) + + +def _client_port(response: httpx.Response) -> int: + stream: Final[object] = response.extensions["network_stream"] # pyright: ignore[reportAny] # httpx types extensions as Any + assert isinstance(stream, httpcore.NetworkStream), stream + return _CLIENT_ADDRESS.validate_python(stream.get_extra_info("client_addr"))[1] + + +def _probe(gateway: Gateway, body: Mapping[str, JsonValue], workers: frozenset[int]) -> tuple[int | None, JsonValue]: + with httpx.Client(base_url=str(gateway.client.base_url), timeout=30, trust_env=False) as client: + response: Final = client.post( + "/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {gateway.key}"} + ) + return _holder(workers, _client_port(response)), JSON_OBJECT.validate_json(response.content).get("input_tokens") + + +def _round( + gateway: Gateway, body: Mapping[str, JsonValue], workers: frozenset[int] +) -> tuple[tuple[int | None, JsonValue], ...]: + def probe(_index: int) -> tuple[int | None, JsonValue]: + return _probe(gateway, body, workers) + + with ThreadPoolExecutor(max_workers=_PROBES_PER_ROUND) as pool: + return tuple(pool.map(probe, range(_PROBES_PER_ROUND))) + + +def _settled_on_every_worker(gateway: Gateway, body: Mapping[str, JsonValue]) -> None: + proxy_port: Final = gateway.client.base_url.port + assert proxy_port is not None, gateway.client.base_url + workers: Final = _proxy_workers(proxy_port) + eventually( + lambda: _round(gateway, body, workers), + lambda observed: ( + frozenset(pid for pid, _ in observed) == workers and all(count == _COUNT for _, count in observed) + ), + seconds=60, + ) + + +def _deploy(gateway: Gateway, scenario: Scenario, provider: _Provider) -> _Deployment: + port: Final = _reserved_port() + model: Final = scenario.model( + model_info=None, model=f"{provider.prefix}/{_MODEL}", api_base=f"http://127.0.0.1:{port}", api_key=_API_KEY + ) + with wire_server(_peer(provider), port=port): + _settled_on_every_worker(gateway, {"model": model, "messages": [USER]}) + return _Deployment(provider, port, model) + + +@pytest.fixture(scope="module") +def deployments() -> Iterator[Mapping[str, _Deployment]]: + with gateway_from_environment() as gateway, gateway.scenario() as scenario: + yield MappingProxyType({chosen.prefix: _deploy(gateway, scenario, chosen) for chosen in (_ANTHROPIC, _AZURE)}) + + +@pytest.fixture +def deployment(provider: _Provider, deployments: Mapping[str, _Deployment]) -> _Deployment: + return deployments[provider.prefix] + + +def _count(gateway: Gateway, body: Mapping[str, JsonValue], key: str | None = None) -> httpx.Response: + return gateway.request("POST", "/v1/messages/count_tokens", body, key=key) + + +def _payload(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text + return JSON_OBJECT.validate_json(response.content) + + +def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int: + response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "false"}) + payload: Final = _payload(response) + total: Final = payload["total_tokens"] + assert payload["tokenizer_type"] not in _PROVIDER_TOKENIZERS, response.text + assert isinstance(total, int) and total > 0 and total != _COUNT, response.text + return total + + +def _bodies(wire: Wire, provider: _Provider) -> tuple[dict[str, JsonValue], ...]: + received: Final = wire.drain() + for request in received: + assert (request.method, request.target) == ("POST", provider.target), request.target + assert request.headers["anthropic-version"] == ANTHROPIC_VERSION, request.headers + assert TOKEN_COUNTING_BETA in request.headers["anthropic-beta"], request.headers + assert request.headers["content-type"] == "application/json", request.headers + assert request.headers["x-api-key"] == _API_KEY, request.headers + assert (request.headers.get("api-key") == _API_KEY) is provider.azure, request.headers + return tuple(JSON_OBJECT.validate_json(request.body) for request in received) + + +def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]: + return tuple( + stack.enter_context(httpx.Client(base_url=base_url, timeout=30, trust_env=False)) for _ in range(count) + ) + + +def _counted_on(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue]: + response: Final = client.post( + "/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {key}"} + ) + return response.status_code, JSON_OBJECT.validate_json(response.content).get("input_tokens") + + +@pytest.mark.parametrize("case", LIFT_CASES.values(), ids=LIFT_CASES.keys()) +def test_messages_count_tokens_lifts_the_leading_system_run( + deployment: _Deployment, gateway: Gateway, case: LiftCase +) -> None: + with _listening(deployment) as peer: + response: Final = _count(gateway, count_request(deployment.model, case)) + assert _payload(response) == {"input_tokens": _COUNT}, response.text + assert _bodies(peer, deployment.provider) == (expected_count_body(_MODEL, case),) + + +def test_anthropic_sdk_count_tokens_lifts_the_leading_system_message(deployment: _Deployment, gateway: Gateway) -> None: + with _listening(deployment) as peer: + counted: Final = anthropic_client(gateway).messages.count_tokens(model=deployment.model, messages=_SDK_MESSAGES) + assert counted.input_tokens == _COUNT, counted + assert _bodies(peer, deployment.provider) == (_STRING_BODY,) + + +def test_async_anthropic_sdk_count_tokens_lifts_the_leading_system_message( + deployment: _Deployment, gateway: Gateway +) -> None: + with _listening(deployment) as peer: + counted: Final = asyncio.run( + async_anthropic_client(gateway).messages.count_tokens(model=deployment.model, messages=_SDK_MESSAGES) + ) + assert counted.input_tokens == _COUNT, counted + assert _bodies(peer, deployment.provider) == (_STRING_BODY,) + + +def test_utils_token_counter_call_endpoint_counts_a_leading_system_through_the_provider( + deployment: _Deployment, gateway: Gateway +) -> None: + with _listening(deployment) as peer: + response: Final = gateway.request( + "POST", + "/utils/token_counter", + count_request(deployment.model, STRING_CASE), + params={"call_endpoint": "true"}, + ) + payload: Final = _payload(response) + expected_tokenizer: Final = deployment.provider.tokenizer + assert (payload["total_tokens"], payload["tokenizer_type"]) == (_COUNT, expected_tokenizer), response.text + assert payload["original_response"] == {"input_tokens": _COUNT}, response.text + assert (payload["request_model"], payload["model_used"]) == (deployment.model, _MODEL), response.text + assert _bodies(peer, deployment.provider) == (_STRING_BODY,) + + +def test_utils_token_counter_local_mode_never_calls_the_peer_for_a_leading_system( + deployment: _Deployment, gateway: Gateway +) -> None: + with _listening(deployment) as peer: + assert _local_count(gateway, count_request(deployment.model, STRING_CASE)) > 0 + assert peer.drain() == () + + +def test_responses_input_tokens_lifts_instructions(deployment: _Deployment, gateway: Gateway) -> None: + with _listening(deployment) as peer: + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": deployment.model, "input": USER_TEXT, "instructions": INSTRUCTION}, + ) + assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _COUNT}, response.text + assert _bodies(peer, deployment.provider) == (_STRING_BODY,) + + +def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_item( + deployment: _Deployment, gateway: Gateway +) -> None: + with _listening(deployment) as peer: + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": deployment.model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION}, + ) + assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _COUNT}, response.text + assert _bodies(peer, deployment.provider) == ( + {"model": _MODEL, "messages": [USER], "system": [LIFTED, {"type": "text", "text": REMINDER}]}, + ) + + +def test_openai_sdk_input_tokens_lifts_instructions(deployment: _Deployment, gateway: Gateway) -> None: + with _listening(deployment) as peer: + counted: Final = openai_client(gateway).responses.input_tokens.count( + model=deployment.model, input=USER_TEXT, instructions=INSTRUCTION + ) + assert counted.input_tokens == _COUNT, counted + assert _bodies(peer, deployment.provider) == (_STRING_BODY,) + + +def test_async_openai_sdk_input_tokens_lifts_instructions(deployment: _Deployment, gateway: Gateway) -> None: + with _listening(deployment) as peer: + counted: Final = asyncio.run( + async_openai_client(gateway).responses.input_tokens.count( + model=deployment.model, input=USER_TEXT, instructions=INSTRUCTION + ) + ) + assert counted.input_tokens == _COUNT, counted + assert _bodies(peer, deployment.provider) == (_STRING_BODY,) + + +def test_messages_count_tokens_forwards_tools_beside_the_lifted_system( + deployment: _Deployment, gateway: Gateway +) -> None: + with _listening(deployment) as peer: + response: Final = _count(gateway, {**count_request(deployment.model, STRING_CASE), "tools": TOOLS}) + assert _payload(response) == {"input_tokens": _COUNT}, response.text + assert _bodies(peer, deployment.provider) == ({**_STRING_BODY, "tools": TOOLS},) + + +def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place( + deployment: _Deployment, gateway: Gateway +) -> None: + with _listening(deployment) as peer: + body: Final[dict[str, JsonValue]] = { + "model": deployment.model, + "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], + } + response: Final = _count(gateway, body) + assert _payload(response) == {"input_tokens": _COUNT}, response.text + assert _bodies(peer, deployment.provider) == ( + {"model": _MODEL, "messages": [USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], "system": [LIFTED]}, + ) + + +def test_messages_count_tokens_falls_back_locally_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 _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [], "system": [LIFTED]},) + + +def test_messages_count_tokens_leaves_the_request_untouched_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 _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [LEADING, USER], "system": 5},) + + +def test_messages_count_tokens_repeated_request_lifts_each_time(deployment: _Deployment, gateway: Gateway) -> None: + with _listening(deployment) as peer: + answers: Final = tuple( + _payload(_count(gateway, count_request(deployment.model, STRING_CASE))) for _ in range(2) + ) + assert answers == ({"input_tokens": _COUNT},) * 2 + assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * 2 + + +def test_messages_count_tokens_answers_a_leading_system_without_content_before_any_peer_call( + 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 peer.drain() == () + follow_up: Final = _count(gateway, count_request(deployment.model, STRING_CASE)) + assert _payload(follow_up) == {"input_tokens": _COUNT}, follow_up.text + assert _bodies(peer, deployment.provider) == (_STRING_BODY,) + + +def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value( + deployment: _Deployment, gateway: Gateway +) -> None: + with _listening(deployment) as peer: + first: Final = json.dumps([USER]) + last: Final = json.dumps([LEADING, USER]) + response: Final = gateway.client.post( + "/v1/messages/count_tokens", + content=f'{{"model": "{deployment.model}", "messages": {first}, "messages": {last}}}', + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + assert _payload(response) == {"input_tokens": _COUNT}, response.text + assert _bodies(peer, deployment.provider) == (_STRING_BODY,) + + +@pytest.mark.parametrize("status", [400, 403, 404, 500, 503]) +def test_messages_count_tokens_falls_back_locally_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 _bodies(peer, deployment.provider) == (_STRING_BODY,) + + +def test_messages_count_tokens_unauthenticated_request_never_reaches_the_peer( + deployment: _Deployment, gateway: Gateway +) -> None: + with _listening(deployment) as peer: + response: Final = _count(gateway, count_request(deployment.model, STRING_CASE), key="sk-not-a-key") + assert response.status_code == 401, response.text + assert peer.drain() == () + + +def test_peer_outage_between_concurrent_waves_falls_back_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) + + 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) + with _listening(deployment) as revived: + assert tuple(pool.map(count, clients)) == ((200, _COUNT),) * len(clients) + assert _bodies(revived, deployment.provider) == (_STRING_BODY,) * len(clients) + + +def test_slow_peer_holds_concurrent_lifted_counts_without_stalling_the_proxy( + deployment: _Deployment, gateway: Gateway +) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + with ExitStack() as stack: + clients: Final = _clients(stack, str(gateway.client.base_url), 6) + peer: Final = stack.enter_context(_listening(deployment, _holding(held, release, 20))) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + stack.callback(release.set) + body: Final = count_request(deployment.model, STRING_CASE) + futures: Final = tuple(pool.submit(_counted_on, client, gateway.key, body) for client in clients) + eventually(held.qsize, lambda size: size == len(clients), seconds=30) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + assert _local_count(gateway, body) > 0 + assert not any(future.done() for future in futures) + release.set() + assert tuple(future.result(timeout=30) for future in futures) == ((200, _COUNT),) * len(clients) + assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * len(clients) + + +def test_chat_completions_on_the_same_deployment_keeps_a_mid_conversation_system_in_place( + deployments: Mapping[str, _Deployment], gateway: Gateway +) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + anthropic_deployment: Final = deployments[_ANTHROPIC.prefix] + + def respond(request: Request) -> Reply: + assert (request.method, request.target) == ("POST", "/v1/messages"), request.target + return _json_reply( + 200, + { + "id": identity, + "type": "message", + "role": "assistant", + "model": _MODEL, + "content": [{"type": "text", "text": "done"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 3}, + }, + ) + + with wire_server(respond, port=anthropic_deployment.port) as peer: + reminder: Final[dict[str, JsonValue]] = { + "role": "system", + "content": [ + {"type": "text", "text": REMINDER, "cache_control": EPHEMERAL}, + {"type": "text", "text": ""}, + IMAGE_PART, + ], + } + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": anthropic_deployment.model, + "max_tokens": 16, + "messages": [LEADING, USER, reminder, ASSISTANT, FOLLOW_UP], + }, + ) + assert response.status_code == 200, response.text + (sent,) = peer.drain() + body: Final = JSON_OBJECT.validate_json(sent.body) + assert body["system"] == [LIFTED], body + assert body["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": USER_TEXT}]}, + {"role": "system", "content": [{"type": "text", "text": REMINDER, "cache_control": EPHEMERAL}]}, + {"role": "assistant", "content": [{"type": "text", "text": "One."}]}, + {"role": "user", "content": [{"type": "text", "text": "Again"}]}, + ], body 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 f9d4f7f91bb..933fc49c149 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -11,7 +11,7 @@ from contextlib import ExitStack from pathlib import Path from queue import SimpleQueue from types import MappingProxyType -from typing import Final +from typing import Final, cast import anthropic import httpx @@ -24,6 +24,25 @@ from integration._support.bedrock_runtime_peer import respond as runtime_generat from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._count_tokens_system_lift import ( + ASSISTANT, + FOLLOW_UP, + INSTRUCTION, + LEADING, + LIFT_CASES, + LIFTED, + MID_SYSTEM, + REMINDER, + STRING_CASE, + USER, + USER_TEXT, + LiftCase, + accepts_count_body, + async_openai_client, + count_request, + expected_count_body, + openai_client, +) from pydantic import JsonValue, TypeAdapter pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @@ -81,6 +100,10 @@ def _mantle_body(**fields: JsonValue) -> dict[str, JsonValue]: _MANTLE_BARE: Final = _mantle_body() +_MANTLE_LIFTED: Final = _mantle_body(system=[LIFTED]) +_LIFT_SDK_MESSAGES: Final = cast( + list[MessageParam], [LEADING, USER] +) # cast-ok: the SDK types reject the role the proxy lifts _MANTLE_FULL: Final = _mantle_body(system=_SYSTEM, tools=_TOOLS) @@ -138,25 +161,8 @@ def _rejecting(status: int) -> Callable[[Request], Reply]: return count -def _anthropic_message(message: JsonValue) -> bool: - return isinstance(message, dict) and message.get("role") in ("user", "assistant") - - -def _anthropic_tool(tool: JsonValue) -> bool: - return isinstance(tool, dict) and isinstance(tool.get("name"), str) and isinstance(tool.get("input_schema"), dict) - - def _strict(request: Request) -> Reply: - body: Final = _JSON_OBJECT.validate_json(request.body) - messages: Final = body.get("messages") - tools: Final = body.get("tools", []) - accepted: Final = ( - isinstance(messages, list) - and all(map(_anthropic_message, messages)) - and isinstance(body.get("system", ""), (str, list)) - and isinstance(tools, list) - and all(map(_anthropic_tool, tools)) - ) + accepted: Final = accepts_count_body(_JSON_OBJECT.validate_json(request.body)) return _mantle_counted(request) if accepted else _rejected(400) @@ -577,7 +583,7 @@ def test_bedrock_passthrough_count_tokens_still_answers_the_runtime_rejection( assert mantle.drain() == () -def test_responses_input_tokens_with_instructions_still_counts_locally( +def test_responses_input_tokens_with_instructions_counts_through_mantle( counting_proxy: OwnedProxy, mantle_port: int ) -> None: gateway: Final = counting_proxy.gateway @@ -592,14 +598,322 @@ def test_responses_input_tokens_with_instructions_still_counts_locally( "/v1/responses/input_tokens", {"model": model, "input": "Count this message", "instructions": "Be terse"}, ) - payload: Final = _payload(response) + assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text assert len(_runtime_count_targets(runtime)) == 1 - (sent,) = _mantle_bodies(mantle) - messages: Final = sent["messages"] - assert isinstance(messages, list) and messages[0] == {"role": "system", "content": "Be terse"}, sent - assert "system" not in sent, sent - local: Final = _local_count(gateway, {"model": model, "messages": messages}) - assert payload == {"object": "response.input_tokens", "input_tokens": local}, response.text + assert _mantle_bodies(mantle) == (_mantle_body(system=[{"type": "text", "text": "Be terse"}]),) + + +@pytest.mark.parametrize("case", LIFT_CASES.values(), ids=LIFT_CASES.keys()) +def test_messages_count_tokens_lifts_the_leading_system_run_for_mantle( + counting_proxy: OwnedProxy, mantle_port: int, case: LiftCase +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, count_request(model, case)) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (expected_count_body(_OPUS_BASE, case),) + + +def test_anthropic_sdk_count_tokens_lifts_the_leading_system_through_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + counted: Final = _anthropic_client(gateway).messages.count_tokens(model=model, messages=_LIFT_SDK_MESSAGES) + assert counted.input_tokens == _MANTLE_COUNT, counted + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) + + +def test_async_anthropic_sdk_count_tokens_lifts_the_leading_system_through_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + counted: Final = asyncio.run( + _async_anthropic_client(gateway).messages.count_tokens(model=model, messages=_LIFT_SDK_MESSAGES) + ) + assert counted.input_tokens == _MANTLE_COUNT, counted + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) + + +def test_utils_token_counter_call_endpoint_counts_a_leading_system_through_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", "/utils/token_counter", count_request(model, STRING_CASE), params={"call_endpoint": "true"} + ) + payload: Final = _payload(response) + assert (payload["total_tokens"], payload["tokenizer_type"]) == (_MANTLE_COUNT, "bedrock_mantle_api") + assert payload["original_response"] == {"input_tokens": _MANTLE_COUNT}, response.text + assert (payload["request_model"], payload["model_used"]) == (model, _OPUS), response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) + + +def test_openai_sdk_input_tokens_lifts_instructions_through_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + counted: Final = openai_client(gateway).responses.input_tokens.count( + model=model, input=USER_TEXT, instructions=INSTRUCTION + ) + assert counted.input_tokens == _MANTLE_COUNT, counted + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) + + +def test_async_openai_sdk_input_tokens_lifts_instructions_through_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + counted: Final = asyncio.run( + async_openai_client(gateway).responses.input_tokens.count( + model=model, input=USER_TEXT, instructions=INSTRUCTION + ) + ) + assert counted.input_tokens == _MANTLE_COUNT, counted + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) + + +def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_item_through_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION}, + ) + assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED, {"type": "text", "text": REMINDER}]),) + + +def test_messages_count_tokens_forwards_tools_beside_the_lifted_system_to_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {**count_request(model, STRING_CASE), "tools": _TOOLS}) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED], tools=_TOOLS),) + + +def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place_for_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + body: Final[dict[str, JsonValue]] = { + "model": model, + "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], + } + response: Final = _count(gateway, body) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == ( + _mantle_body(messages=[USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], system=[LIFTED]), + ) + + +def test_messages_count_tokens_falls_back_locally_when_every_message_is_system_for_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + 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 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( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + 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 len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(messages=[LEADING, USER], system=5),) + + +def test_messages_count_tokens_repeated_request_lifts_the_leading_system_each_time_for_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + answers: Final = tuple(_payload(_count(gateway, count_request(model, STRING_CASE))) for _ in range(2)) + assert answers == ({"input_tokens": _MANTLE_COUNT},) * 2 + assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 + assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * 2 + + +def test_messages_count_tokens_answers_a_leading_system_without_content_before_any_mantle_call( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + 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 mantle.drain() == () + follow_up: Final = _count(gateway, count_request(model, STRING_CASE)) + assert _payload(follow_up) == {"input_tokens": _MANTLE_COUNT}, follow_up.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) + + +def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value_for_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_strict), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + first: Final = json.dumps([USER]) + last: Final = json.dumps([LEADING, USER]) + response: Final = gateway.client.post( + "/v1/messages/count_tokens", + content=f'{{"model": "{model}", "messages": {first}, "messages": {last}}}', + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) + + +def test_disabled_token_counter_counts_a_leading_system_through_mantle( + gateway: Gateway, mantle_port: int, tmp_path: Path +) -> None: + with ExitStack() as stack: + runtime: Final = stack.enter_context(wire_server(_runtime)) + config: Final = _owned_config( + tmp_path / "disabled-token-counter-lift.yaml", runtime.url, {"disable_token_counter": True} + ) + owned: Final = stack.enter_context( + owned_proxy_process( + gateway, + tmp_path, + _mantle_environment(mantle_port), + config=config, + workers=2, + remove_environment=_INHERITED_BEARER, + ) + ) + with wire_server(_mantle(_strict), port=mantle_port) as strict: + counted: Final = _count(owned.gateway, count_request(_OWNED_OPUS, STRING_CASE)) + assert _payload(counted) == {"input_tokens": _MANTLE_COUNT}, counted.text + assert _mantle_bodies(strict) == (_MANTLE_LIFTED,) + assert len(_runtime_count_targets(runtime)) == 1 + + +def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ExitStack() as stack: + clients: Final = _clients(stack, str(gateway.client.base_url), 8) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + runtime: Final = stack.enter_context(wire_server(_runtime)) + 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) + + 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) + 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) + assert len(_runtime_count_targets(runtime)) == 3 * len(clients) @pytest.mark.parametrize("status", [400, 403, 404, 500, 503]) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py index d1e23a17747..9e3704a9df0 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py @@ -11,6 +11,7 @@ import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( CONVERTED_SYSTEM_NOTE, + anthropic_system_blocks, place_mid_conversation_system, split_leading_system_run, ) @@ -390,3 +391,23 @@ def test_flagged_placement_keeps_a_system_before_an_assistant_turn_whose_empty_t ) assert _roles(placed) == ["user", "system", "assistant", "user"] + + +def test_anthropic_system_blocks_keeps_text_parts_with_their_cache_control_and_drops_the_rest(): + run = [ + {"role": "system", "content": "one", "cache_control": {"type": "ephemeral", "ttl": "1h"}}, + { + "role": "system", + "content": [ + {"type": "text", "text": "two", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": ""}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}}, + ], + }, + {"role": "system", "content": ""}, + ] + + assert anthropic_system_blocks(run) == ( + {"type": "text", "text": "one", "cache_control": {"type": "ephemeral", "ttl": "1h"}}, + {"type": "text", "text": "two", "cache_control": {"type": "ephemeral"}}, + ) diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index 18d57e71d28..fea65a9c821 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -267,3 +267,115 @@ async def test_remote_image_fetch_keeps_counting_handler_event_loop_responsive( ]}], } assert messages == original + + +@pytest.mark.parametrize( + "config_type", (AnthropicCountTokensConfig, AzureAIAnthropicCountTokensConfig) +) +def test_count_lifts_the_leading_system_run_into_system( + config_type: type[AnthropicCountTokensConfig], +) -> None: + """A Responses ``instructions`` arrives as a leading system-role message. count_tokens answers 400 + on that role at the head of ``messages`` and only takes the initial prompt in ``system``, so the + leading run moves there with its cache_control, empty text dropped, and a later reminder stays.""" + cache_control: Final[dict[str, JsonValue]] = {"type": "ephemeral"} + messages: Final[list[dict[str, JsonValue]]] = [ + {"role": "system", "content": "Be terse", "cache_control": cache_control}, + {"role": "system", "content": [{"type": "text", "text": "Answer in French"}, {"type": "text", "text": ""}]}, + {"role": "user", "content": "Hello, how are you?"}, + {"role": "system", "content": "later reminder"}, + {"role": "assistant", "content": "Bonjour."}, + ] + original: Final = deepcopy(messages) + result: Final = config_type().transform_request_to_count_tokens(model="claude-opus-5-5", messages=messages) + + assert result == { + "model": "claude-opus-5-5", + "system": [ + {"type": "text", "text": "Be terse", "cache_control": cache_control}, + {"type": "text", "text": "Answer in French"}, + ], + "messages": [ + {"role": "user", "content": "Hello, how are you?"}, + {"role": "system", "content": "later reminder"}, + {"role": "assistant", "content": "Bonjour."}, + ], + } + assert messages == original + + +@pytest.mark.parametrize( + ("system", "expected_system"), + ( + (None, [{"type": "text", "text": "Be terse"}]), + ("", [{"type": "text", "text": "Be terse"}]), + ("Answer in French", [{"type": "text", "text": "Answer in French"}, {"type": "text", "text": "Be terse"}]), + ( + [{"type": "text", "text": "Answer in French", "cache_control": {"type": "ephemeral"}}], + [ + {"type": "text", "text": "Answer in French", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "Be terse"}, + ], + ), + ), + ids=["absent", "empty", "string", "blocks"], +) +def test_count_keeps_the_callers_system_ahead_of_the_lifted_run( + system: JsonValue, expected_system: list[dict[str, JsonValue]] +) -> None: + result: Final = AnthropicCountTokensConfig().transform_request_to_count_tokens( + model="claude-opus-5-5", + messages=[{"role": "system", "content": "Be terse"}, {"role": "user", "content": "hi"}], + system=system, + ) + + assert result == { + "model": "claude-opus-5-5", + "system": expected_system, + "messages": [{"role": "user", "content": "hi"}], + } + + +def test_count_leaves_a_non_text_system_and_its_messages_as_sent() -> None: + """A malformed ``system`` is the provider's to reject, so nothing is rearranged around it.""" + messages: Final[list[dict[str, JsonValue]]] = [ + {"role": "system", "content": "Be terse"}, + {"role": "user", "content": "hi"}, + ] + result: Final = AnthropicCountTokensConfig().transform_request_to_count_tokens( + model="claude-opus-5-5", messages=messages, system=5 + ) + + assert result == {"model": "claude-opus-5-5", "system": 5, "messages": messages} + + +def test_count_drops_a_leading_system_message_without_text() -> None: + result: Final = AnthropicCountTokensConfig().transform_request_to_count_tokens( + model="claude-opus-5-5", + messages=[{"role": "system", "content": ""}, {"role": "user", "content": "hi"}], + ) + + assert result == {"model": "claude-opus-5-5", "messages": [{"role": "user", "content": "hi"}]} + + +@pytest.mark.asyncio +async def test_handler_sends_the_leading_system_run_as_system_not_as_a_message(httpx_transport_clients): + """The wire body is what the provider judges: ``system`` carries the prompt and no message has + ``role: "system"``, so a Responses ``instructions`` is counted by Anthropic instead of 400ing.""" + with respx.mock: + route = respx.post(url__startswith="https://gateway.example").mock( + return_value=httpx.Response(200, json={"input_tokens": 21}) + ) + result = await AnthropicCountTokensHandler().handle_count_tokens_request( + model="claude-opus-5-5", + messages=[{"role": "system", "content": "Be terse"}, {"role": "user", "content": "Hello, how are you?"}], + api_key="sk-ant-api03-test-key", + api_base="https://gateway.example", + ) + + assert result == {"input_tokens": 21} + assert TypeAdapter(dict[str, JsonValue]).validate_json(route.calls.last.request.content) == { + "model": "claude-opus-5-5", + "system": [{"type": "text", "text": "Be terse"}], + "messages": [{"role": "user", "content": "Hello, how are you?"}], + } From c703543b302309ed8589d3ae5e51c6d5897674d6 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 19 Oct 2026 10:18:05 -0700 Subject: [PATCH 11/36] fix(proxy): raise on provider count_tokens errors instead of falling back to tiktoken (#46121) Backport of #46121 to stable/1.104.x. Cherry-picked from 719c44db3f2ec5a13d51fbf29524e001720306be (main). Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yuneng-jiang --- litellm/proxy/proxy_server.py | 17 ++--- ...anthropic_count_tokens_system_lift_wire.py | 31 ++++---- .../test_bedrock_mantle_count_tokens_wire.py | 71 ++++++++++--------- .../test_vertex_partner_count_tokens_wire.py | 49 +++++++------ tests/unit/proxy/test_proxy_token_counter.py | 55 +++++++------- 5 files changed, 110 insertions(+), 113 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 47fe4883f3c..76e01a874a2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13643,19 +13643,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 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 13091ccd14a..a4b5c15ad54 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 @@ -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) 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 933fc49c149..276265a521c 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -523,7 +523,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 @@ -534,7 +534,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() == () @@ -552,17 +553,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() == () @@ -771,7 +774,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 @@ -782,14 +785,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 @@ -799,9 +802,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),) @@ -822,7 +825,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 @@ -833,9 +836,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 @@ -890,7 +893,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 @@ -901,7 +904,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) @@ -909,7 +911,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) @@ -917,7 +921,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 @@ -928,7 +932,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,) @@ -938,7 +943,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 @@ -949,21 +954,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 @@ -984,18 +990,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 @@ -1006,7 +1012,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),) @@ -1175,7 +1182,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 @@ -1186,7 +1193,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) @@ -1198,7 +1204,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) diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index 5df8dd6e32e..ae9b8434d11 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -345,15 +345,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: @@ -366,7 +365,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( @@ -374,11 +373,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: @@ -389,26 +387,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( @@ -422,17 +420,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( @@ -477,13 +476,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: @@ -579,7 +578,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))) @@ -588,7 +587,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) @@ -596,7 +594,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) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 4bd48e1effc..e4832a7c23b 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -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 = ( From fe9e6381a350e827f31ad58f50add15ca691f26c Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Mon, 19 Oct 2026 10:23:45 -0700 Subject: [PATCH 12/36] fix(count_tokens): count system and tools on Vertex Gemini, fall back to Mantle on Bedrock 404 (#46430) Backport of #46430 to stable/1.104.x. Cherry-picked from 08a24d391af6ba0fbd4e99313857fbb05754ddf2 (main). Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../count_tokens/bedrock_token_counter.py | 5 +- litellm/llms/gemini/common_utils.py | 11 + litellm/llms/gemini/count_tokens/handler.py | 44 +- .../responses/count_tokens/transformation.py | 13 +- litellm/llms/vertex_ai/common_utils.py | 32 +- .../llms/vertex_ai/count_tokens/context.py | 144 ++++++ .../llms/vertex_ai/count_tokens/handler.py | 32 +- litellm/proxy/google_endpoints/endpoints.py | 96 +++- .../test_bedrock_mantle_count_tokens_wire.py | 134 +++++ .../test_openai_count_tokens_tools_wire.py | 161 ++++++ .../test_vertex_gemini_count_tokens_wire.py | 462 ++++++++++++++++++ .../proxy/google_endpoints/test_endpoints.py | 80 ++- .../test_bedrock_token_counter.py | 38 +- .../llms/gemini/test_gemini_common_utils.py | 146 +++++- ...test_openai_count_tokens_transformation.py | 70 ++- .../llms/vertex_ai/count_tokens/__init__.py | 0 .../vertex_ai/count_tokens/test_context.py | 196 ++++++++ .../vertex_ai/test_vertex_ai_common_utils.py | 149 +++--- 18 files changed, 1691 insertions(+), 122 deletions(-) create mode 100644 litellm/llms/vertex_ai/count_tokens/context.py create mode 100644 tests/integration/providers/test_openai_count_tokens_tools_wire.py create mode 100644 tests/integration/providers/test_vertex_gemini_count_tokens_wire.py create mode 100644 tests/unit/llms/vertex_ai/count_tokens/__init__.py create mode 100644 tests/unit/llms/vertex_ai/count_tokens/test_context.py diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index ed61707f99e..3084b0fde07 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -20,6 +20,7 @@ from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCo RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" +_RUNTIME_FALLBACK_STATUSES: Final = frozenset({400, 404}) class _CountTokensReply(LiteLLMPydanticObjectBase): @@ -46,14 +47,14 @@ CountTokensOutcome = CountedTokens | CountTokensFailure def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool: return ( isinstance(outcome, CountTokensFailure) - and outcome.status_code == 400 + and outcome.status_code in _RUNTIME_FALLBACK_STATUSES and is_mantle_claude_model(resolved_model) ) class BedrockTokenCounter(BaseTokenCounter): """Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens - on bedrock-mantle for the Claude models bedrock-runtime answers 400 for""" + on bedrock-mantle for Claude models bedrock-runtime answers 400 or 404 for""" def __init__( self, diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 78e6e6aaf82..2c3b163496e 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -495,12 +495,23 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): import copy from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter + from litellm.llms.vertex_ai.count_tokens.context import ( + VertexGeminiCountContextError, + count_context_error_response, + vertex_gemini_count_context, + ) + + count_context: Final = vertex_gemini_count_context(system=system, tools=tools) + if isinstance(count_context, VertexGeminiCountContextError): + return count_context_error_response(count_context, request_model, model_to_use) deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) count_tokens_params: Final = { "model": model_to_use, "contents": contents, + "system_instruction": count_context.system_instruction, + "tools": count_context.tools, } count_tokens_params_request.update(count_tokens_params) result: Final = await GoogleAIStudioTokenCounter().acount_tokens( diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py index c2f0ef473ae..e11a401a5f9 100644 --- a/litellm/llms/gemini/count_tokens/handler.py +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -1,9 +1,10 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx import litellm -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.types.utils import LlmProviders if TYPE_CHECKING: @@ -28,12 +29,12 @@ class GoogleAIStudioTokenCounter: """ import copy - from google.genai.types import FunctionResponse - # Handle None or empty contents if not contents: return contents + from google.genai.types import FunctionResponse + cleaned_contents: Final = copy.deepcopy(contents) for content in cleaned_contents: @@ -54,6 +55,28 @@ class GoogleAIStudioTokenCounter: base_url: Final = api_base or "https://generativelanguage.googleapis.com" return f"{base_url}/v1beta/models/{model}:countTokens" + def __init__(self, client: AsyncHTTPHandler | None = None) -> None: + super().__init__() + self._client: Final = client + + def _build_count_tokens_request_body( + self, + model: str, + contents: object, + system_instruction: Mapping[str, object] | None, + tools: Sequence[Mapping[str, object]] | None, + ) -> Mapping[str, object]: + if system_instruction is None and not tools: + return {"contents": contents} + + generate_content_request: Final = { + "model": model if model.startswith("models/") else f"models/{model}", + "contents": contents, + **({"systemInstruction": dict(system_instruction)} if system_instruction is not None else {}), + **({"tools": [dict(tool) for tool in tools]} if tools else {}), + } + return {"generateContentRequest": generate_content_request} + async def validate_environment( self, api_base: str | None = None, @@ -84,6 +107,8 @@ class GoogleAIStudioTokenCounter: api_key: str | None = None, api_base: str | None = None, timeout: float | httpx.Timeout | None = None, + system_instruction: Mapping[str, object] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, **kwargs: object, ) -> dict[str, Any]: """ @@ -130,14 +155,19 @@ class GoogleAIStudioTokenCounter: # Prepare request body - clean up contents to remove unsupported fields cleaned_contents: Final = self._clean_contents_for_gemini_api(contents) - request_body: Final = {"contents": cleaned_contents} + request_body: Final = self._build_count_tokens_request_body( + model=model, + contents=cleaned_contents, + system_instruction=system_instruction, + tools=tools, + ) - async_httpx_client: Final = get_async_httpx_client( - llm_provider=LlmProviders.GEMINI, + async_httpx_client: Final = ( + self._client if self._client is not None else get_async_httpx_client(llm_provider=LlmProviders.GEMINI) ) try: - response: Final = await async_httpx_client.post(url=url, headers=headers, json=request_body) + response: Final = await async_httpx_client.post(url=url, headers=headers, json=dict(request_body)) # Check for HTTP errors response.raise_for_status() diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 9f596505f91..fa46fdc69ef 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -157,9 +157,10 @@ class OpenAICountTokensConfig: tools: list[dict[str, Any]], ) -> list[dict[str, object]]: """ - Transform OpenAI chat tools format to Responses API tools format. + Transform OpenAI chat and Anthropic tools to Responses API tools format. Chat format: {"type": "function", "function": {"name": "...", "parameters": {...}}} + Anthropic format: {"name": "...", "input_schema": {...}} Responses format: {"type": "function", "name": "...", "parameters": {...}} """ transformed: Final[list[dict[str, object]]] = [] @@ -175,6 +176,16 @@ class OpenAICountTokensConfig: if "strict" in func: item["strict"] = func["strict"] transformed.append(item) + elif "input_schema" in tool and "name" in tool and tool.get("type") in (None, "custom"): + transformed.append( + { + "type": "function", + "name": tool["name"], + "description": tool.get("description", ""), + "parameters": tool["input_schema"], + "strict": bool(tool.get("strict")), + } + ) else: # Pass through non-function tools (e.g., web_search, file_search) transformed.append(tool) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index ead96da5cea..8643ae31e29 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1308,11 +1308,20 @@ class VertexAITokenCounter(BaseTokenCounter): original_response=result, ) else: + from litellm.llms.vertex_ai.count_tokens.context import ( + VertexGeminiCountContextError, + count_context_error_response, + vertex_gemini_count_context, + ) from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter from litellm.llms.vertex_ai.gemini.transformation import ( _gemini_convert_messages_with_history, # pyright: ignore[reportPrivateUsage] # shared helper already used by gemini/chat, context_caching, and vertex_and_google_ai_studio_gemini ) + count_context: Final = vertex_gemini_count_context(system=system, tools=tools) + if isinstance(count_context, VertexGeminiCountContextError): + return count_context_error_response(count_context, request_model, model_to_use) + resolved_contents: Final = ( contents if contents is not None @@ -1324,18 +1333,31 @@ class VertexAITokenCounter(BaseTokenCounter): count_tokens_params: Final = { "model": model_to_use, "contents": resolved_contents, + "system_instruction": count_context.system_instruction, + "tools": count_context.tools, } count_tokens_params_request.update(count_tokens_params) - result = await VertexAITokenCounter().acount_tokens( - **count_tokens_params_request, - ) + try: + result = await VertexAITokenCounter().acount_tokens( + **count_tokens_params_request, + ) + except litellm.APIError as error: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="gemini", + error=True, + error_message=error.message, + status_code=error.status_code, + ) - if result is not None and "totalTokens" in result: + if "totalTokens" in result: return TokenCountResponse( total_tokens=result["totalTokens"], request_model=request_model, model_used=model_to_use, - tokenizer_type=result.get("tokenizer_used", ""), + tokenizer_type=result.get("tokenizer_used", "gemini"), original_response=result, ) diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py new file mode 100644 index 00000000000..24ffb406a47 --- /dev/null +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -0,0 +1,144 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.types.llms.anthropic import AnthropicMessagesTextParam +from litellm.types.utils import TokenCountResponse + + +@dataclass(frozen=True, slots=True) +class VertexGeminiCountContext: + system_instruction: Mapping[str, object] | None + tools: tuple[Mapping[str, object], ...] | None + + +@dataclass(frozen=True, slots=True) +class VertexGeminiCountContextError: + error_message: str + + +_SYSTEM_BLOCKS: Final = TypeAdapter(tuple[AnthropicMessagesTextParam, ...]) +_TOOLS: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) +_VERTEX_TOOL_KEYS: Final = MappingProxyType( + { + "function_declarations": "functionDeclarations", + "url_context": "urlContext", + "code_execution": "codeExecution", + } +) + + +def _system_parts( + system: object | None, +) -> tuple[Mapping[str, JsonValue], ...] | VertexGeminiCountContextError: + if system is None: + return () + if isinstance(system, str): + return ({"text": system},) if system else () + try: + blocks: Final = _SYSTEM_BLOCKS.validate_python(system) + except ValidationError: + return VertexGeminiCountContextError(error_message="Gemini system instructions only support text blocks.") + return tuple({"text": block["text"]} for block in blocks) + + +_INVALID_TOOLS: Final = VertexGeminiCountContextError( + error_message="tools must be a list of objects, and each function tool needs a JSON object schema." +) + + +def _has_object_schema(tool: Mapping[str, JsonValue]) -> bool: + function: Final = tool.get("function") + schema: Final = ( + function.get("parameters") if isinstance(function, dict) else tool.get("input_schema", tool.get("parameters")) + ) + return schema is None or isinstance(schema, dict) + + +def _is_function_tool(tool: Mapping[str, JsonValue]) -> bool: + return isinstance(tool.get("name"), str) or "input_schema" in tool or "function" in tool + + +def _is_responses_function_tool(tool: Mapping[str, JsonValue]) -> bool: + return tool.get("type") == "function" and "function" not in tool and "input_schema" not in tool + + +def _chat_tool(tool: Mapping[str, JsonValue]) -> tuple[Mapping[str, object], ...]: + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + if _is_responses_function_tool(tool): + chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + [dict(tool)] # pyright: ignore[reportArgumentType] # validated JSON dict in the Responses function tool shape + ) + return tuple(chat_tools) + anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai( + [dict(tool)] # pyright: ignore[reportArgumentType] # validated JSON dict in an Anthropic or chat tool shape + ) + return tuple(anthropic_tools) + + +def _map_function_tools( + tools: object | None, +) -> tuple[Mapping[str, object], ...] | None | VertexGeminiCountContextError: + try: + tool_values: Final = _TOOLS.validate_python(tools or ()) + except ValidationError: + return _INVALID_TOOLS + function_tools: Final = tuple(tool for tool in tool_values if _is_function_tool(tool)) + if not function_tools: + return None + if not all(_has_object_schema(tool) for tool in function_tools): + return _INVALID_TOOLS + + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + chat_tools: Final = tuple( + chat_tool for tool in function_tools for chat_tool in _chat_tool(tool) + ) # comprehension-ok: flattens each tool's translation + try: + mapped_tools: Final = VertexGeminiConfig()._map_function( # pyright: ignore[reportPrivateUsage] # the tool mapping inference uses + value=[dict(tool) for tool in chat_tools], + optional_params={}, + ) + except (AttributeError, TypeError, ValueError): + return _INVALID_TOOLS + return tuple( + MappingProxyType({_VERTEX_TOOL_KEYS.get(key, key): value for key, value in tool.items()}) + for tool in mapped_tools + ) + + +def vertex_gemini_count_context( + system: object | None, + tools: object | None, +) -> VertexGeminiCountContext | VertexGeminiCountContextError: + system_parts: Final = _system_parts(system) + if isinstance(system_parts, VertexGeminiCountContextError): + return system_parts + function_tools: Final = _map_function_tools(tools) + if isinstance(function_tools, VertexGeminiCountContextError): + return function_tools + return VertexGeminiCountContext( + system_instruction=MappingProxyType({"parts": system_parts}) if system_parts else None, + tools=function_tools, + ) + + +def count_context_error_response( + error: VertexGeminiCountContextError, request_model: str, model_used: str +) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_used, + tokenizer_type="gemini", + error=True, + error_message=error.error_message, + status_code=400, + ) diff --git a/litellm/llms/vertex_ai/count_tokens/handler.py b/litellm/llms/vertex_ai/count_tokens/handler.py index 47a08ff054d..03970099f2e 100644 --- a/litellm/llms/vertex_ai/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/count_tokens/handler.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping, Sequence from typing import Any, Final from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter @@ -5,6 +6,19 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase): + def _build_count_tokens_request_body( + self, + model: str, + contents: object, + system_instruction: Mapping[str, object] | None, + tools: Sequence[Mapping[str, object]] | None, + ) -> Mapping[str, object]: + return { + "contents": contents, + **({"systemInstruction": dict(system_instruction)} if system_instruction else {}), + **({"tools": [dict(tool) for tool in tools]} if tools else {}), + } + async def validate_environment( self, api_base: str | None = None, @@ -16,29 +30,29 @@ class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase): """ Returns a Tuple of headers and url for the Vertex AI countTokens endpoint. """ - litellm_params = litellm_params or {} - vertex_credentials: Final = self.get_vertex_ai_credentials(litellm_params=litellm_params) - vertex_project = self.get_vertex_ai_project(litellm_params=litellm_params) - vertex_location: Final = self.get_vertex_ai_location(litellm_params=litellm_params) - _auth_header, vertex_project = await self._ensure_access_token_async( + resolved_litellm_params: Final = litellm_params or {} + vertex_credentials: Final = self.get_vertex_ai_credentials(litellm_params=resolved_litellm_params) + vertex_project: Final = self.get_vertex_ai_project(litellm_params=resolved_litellm_params) + vertex_location: Final = self.get_vertex_ai_location(litellm_params=resolved_litellm_params) + _auth_header, resolved_vertex_project = await self._ensure_access_token_async( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", ) - auth_header, api_base = self._get_token_and_url( + auth_header, url = self._get_token_and_url( model=model, gemini_api_key=None, auth_header=_auth_header, - vertex_project=vertex_project, + vertex_project=resolved_vertex_project, vertex_location=vertex_location, vertex_credentials=vertex_credentials, stream=False, custom_llm_provider="vertex_ai", - api_base=None, + api_base=api_base, mode="count_tokens", ) headers = { "Authorization": f"Bearer {auth_header}", } - return headers, api_base + return headers, url diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 4b115686773..864ca2f9432 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,19 +1,96 @@ +from collections.abc import Mapping +from itertools import chain from typing import Final -from fastapi import APIRouter, Depends, Request, Response +from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import ORJSONResponse +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types from litellm.proxy._types import * +from litellm.proxy._types import TokenCountRequest from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.types.llms.vertex_ai import TokenCountDetailsResponse + +class _GoogleCountTokensTextPart(BaseModel): + model_config = ConfigDict(frozen=True) + text: str + + +class _GoogleCountTokensSystemInstruction(BaseModel): + model_config = ConfigDict(frozen=True) + parts: tuple[_GoogleCountTokensTextPart, ...] + + +class _GoogleCountTokensFunctionDeclaration(BaseModel): + model_config = ConfigDict(frozen=True, populate_by_name=True) + name: str + description: str | None = None + parameters: JsonValue | None = None + parameters_json_schema: JsonValue | None = Field(default=None, alias="parametersJsonSchema") + + +class _GoogleCountTokensTool(BaseModel): + model_config = ConfigDict(frozen=True, populate_by_name=True) + function_declarations: tuple[_GoogleCountTokensFunctionDeclaration, ...] = Field( + default=(), alias="functionDeclarations" + ) + + +class _GoogleCountTokensRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True) + contents: tuple[Mapping[str, JsonValue], ...] = () + system_instruction: _GoogleCountTokensSystemInstruction | None = Field(default=None, alias="systemInstruction") + tools: tuple[_GoogleCountTokensTool, ...] = () + + router: Final = APIRouter( tags=["google genai endpoints"], ) +def _input_schema(declaration: _GoogleCountTokensFunctionDeclaration) -> object: + if declaration.parameters_json_schema is not None: + return declaration.parameters_json_schema + if declaration.parameters is not None: + return normalize_json_schema_types(declaration.parameters) + return {"type": "object", "properties": {}} + + +def _google_count_tokens_request(model_name: str, data: object) -> TokenCountRequest: + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + validated_data: Final = TypeAdapter(_GoogleCountTokensRequest).validate_python(data) + contents: Final = [dict(content) for content in validated_data.contents] + messages: Final = ( + GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", []) + ) + system: Final = ( + [{"type": "text", "text": part.text} for part in validated_data.system_instruction.parts] + if validated_data.system_instruction is not None + else None + ) + declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools) + tools: Final = [ + { + "name": declaration.name, + **({"description": declaration.description} if declaration.description is not None else {}), + "input_schema": _input_schema(declaration), + } + for declaration in declarations + ] + return TokenCountRequest( + model=model_name, + contents=contents, + messages=messages, + system=system, + tools=tools or None, + ) + + @router.post( "/v1beta/models/{model_name:path}:generateContent", dependencies=[Depends(user_api_key_auth)], @@ -165,23 +242,14 @@ async def google_count_tokens(request: Request, model_name: str): } ``` """ - from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.proxy_server import token_counter as internal_token_counter data: Final = await _read_request_body(request=request) - contents: Final = data.get("contents", []) - # Create TokenCountRequest for the internal endpoint - from litellm.proxy._types import TokenCountRequest - - # Translate contents to openai format messages using the adapter - messages = GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", []) - - token_request: Final = TokenCountRequest( - model=model_name, - contents=contents, - messages=messages, # compatibility when use openai-like endpoint - ) + try: + token_request: Final = _google_count_tokens_request(model_name, data) + except ValidationError as error: + raise HTTPException(status_code=400, detail="Invalid Google countTokens request body") from error # Call the internal token counter function with direct request flag set to False token_response: Final = await internal_token_counter( 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 276265a521c..bc70c5d3b21 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -457,6 +457,140 @@ def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock assert mantle.drain() == () +def test_runtime_404_on_claude_falls_back_to_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + system: Final = f"{_SYSTEM} status=404" + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) + + +def test_runtime_404_and_mantle_404_surface_provider_error(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + system: Final = f"{_SYSTEM} status=404" + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_rejecting(404)), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) + assert response.status_code == 404, response.text + assert _REJECTION in response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) + + +def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url, _NOVA) + body: Final = {**_bare(model), "messages": [{"role": "user", "content": "status=404"}]} + response: Final = _count(gateway, body) + assert response.status_code == 404, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert mantle.drain() == () + + +def test_google_count_tokens_forwards_system_and_tools_to_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + system_text: Final = "Use the provided tool" + tool_schema: Final = {"type": "object", "properties": {"location": {"type": "string"}}} + body: Final[dict[str, JsonValue]] = { + "contents": [{"role": "user", "parts": [{"text": "status=404"}]}], + "systemInstruction": {"parts": [{"text": system_text}]}, + "tools": [ + { + "functionDeclarations": [ + {"name": "get_weather", "description": "Look up weather", "parameters": tool_schema} + ] + } + ], + } + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) + assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == ( + _mantle_body( + messages=[{"role": "user", "content": "status=404"}], + system=[{"type": "text", "text": system_text}], + tools=[ + { + "name": "get_weather", + "description": "Look up weather", + "input_schema": tool_schema, + } + ], + ), + ) + + +def test_google_count_tokens_normalizes_uppercase_parameters_for_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + body: Final[dict[str, JsonValue]] = { + "contents": [{"role": "user", "parts": [{"text": "status=404"}]}], + "tools": [ + { + "functionDeclarations": [ + { + "name": "get_weather", + "parameters": { + "type": "OBJECT", + "properties": {"location": {"type": "STRING"}}, + "required": ["location"], + }, + } + ] + } + ], + } + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) + assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == ( + _mantle_body( + messages=[{"role": "user", "content": "status=404"}], + tools=[ + { + "name": "get_weather", + "input_schema": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + } + ], + ), + ) + + def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway with ( diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py new file mode 100644 index 00000000000..bfae84e333d --- /dev/null +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -0,0 +1,161 @@ +import json +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.client import Scenario, gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_COUNT: Final = 4242 +_TOOL_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {}} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@pytest.fixture(scope="module") +def counting_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process(gateway, tmp_path_factory.mktemp("openai-count-tools"), {}, workers=1) as owned, + ): + yield owned + + +def _deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model( + model="openai/gpt-4o-mini", + api_base=f"{api_base}/v1", + api_key="sk-integration-openai", + ) + + +def _count_response(_request: Request) -> Reply: + return Reply(body=json.dumps({"input_tokens": _COUNT}).encode()) + + +def _count_body(wire: Wire) -> dict[str, JsonValue]: + requests: Final = wire.drain() + assert len(requests) == 1, requests + request: Final = requests[0] + assert request.method == "POST" + assert request.target == "/v1/responses/input_tokens", request.target + return _JSON_OBJECT.validate_json(request.body) + + +def test_openai_count_tokens_uses_deployment_api_base(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": [{"role": "user", "content": "hello"}]}, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT + assert "tools" not in _count_body(wire) + + +def test_google_count_tokens_maps_anthropic_tools_to_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + { + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "tools": [ + { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + } + ] + } + ], + }, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["totalTokens"] == _COUNT + assert _count_body(wire)["tools"] == [ + { + "type": "function", + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + "strict": False, + } + ] + + +def test_messages_count_tokens_maps_anthropic_tools_to_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + { + "model": model, + "messages": [{"role": "user", "content": "hello"}], + "tools": [ + { + "name": "lookup", + "description": "Find a record", + "input_schema": _TOOL_SCHEMA, + } + ], + }, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT + assert _count_body(wire)["tools"] == [ + { + "type": "function", + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + "strict": False, + } + ] + + +def test_utils_token_counter_keeps_chat_tool_conversion(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/utils/token_counter", + { + "model": model, + "messages": [{"role": "user", "content": "hello"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + }, + } + ], + }, + params={"call_endpoint": "true"}, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["total_tokens"] == _COUNT + assert _count_body(wire)["tools"] == [ + { + "type": "function", + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + } + ] diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py new file mode 100644 index 00000000000..1a148b19f8d --- /dev/null +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -0,0 +1,462 @@ +import asyncio +import json +import threading +from collections.abc import Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from types import MappingProxyType +from typing import Final +from uuid import uuid4 + +import anthropic +import httpx +import pytest +from integration._support.client import Gateway, Scenario +from integration._support.process import graceful_stop_seconds +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_MODEL: Final = "gemini-4-argon" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_COUNT: Final = 777 +_UNIT: Final = " red blue green yellow orange purple black white silver gold" +_USER: Final = f"{_UNIT * 100}\nReturn exactly OK." +_MESSAGES: Final[JsonValue] = [{"role": "user", "content": _USER}] +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) +_REPLY: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "OK"}]}, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 1, "totalTokenCount": 11}, +} + + +class _GeminiPeer: + def __init__(self) -> None: + self.status: int = 200 + self.requests: tuple[Request, ...] = () + self._lock: Final = threading.Lock() + + def set_status(self, status: int) -> None: + self.status = status + + def __call__(self, request: Request) -> Reply: + with self._lock: + self.requests = (*self.requests, request) + status: Final = self.status + if request.target.endswith(":countTokens"): + return Reply(status=status, body=json.dumps({"totalTokens": _COUNT}).encode()) + if request.target.endswith(":generateContent"): + return Reply(body=json.dumps(_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": request.target}).encode()) + + def drain(self) -> tuple[Request, ...]: + with self._lock: + requests: Final = self.requests + self.requests = () + return requests + + +def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str: + return scenario.model( + model=f"vertex_ai/{_MODEL}", + api_base=api_base, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + ) + + +def _messages(model: str, system: JsonValue | None, tools: JsonValue | None) -> Mapping[str, JsonValue]: + return { + "model": model, + "messages": _MESSAGES, + **({"system": system} if system is not None else {}), + **({"tools": tools} if tools is not None else {}), + } + + +def _json_body(request: Request) -> Mapping[str, JsonValue]: + return _JSON_OBJECT.validate_json(request.body) + + +def _count_requests(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if request.target.endswith(":countTokens")) + + +def _count_tokens_tools(generated_tools: JsonValue) -> JsonValue: + return _JSON_LIST.validate_python( + [ + { + { + "function_declarations": "functionDeclarations", + "url_context": "urlContext", + "code_execution": "codeExecution", + }.get(key, key): value + for key, value in tool.items() + } + for tool in TypeAdapter(list[Mapping[str, JsonValue]]).validate_python(generated_tools) + ] + ) + + +def _count_tokens_request_from_generate_content( + generated_body: Mapping[str, JsonValue], +) -> Mapping[str, JsonValue]: + system_instruction: Final = generated_body.get("systemInstruction", generated_body.get("system_instruction")) + generated_tools: Final = generated_body.get("tools") + return MappingProxyType( + { + "contents": generated_body["contents"], + **({"systemInstruction": system_instruction} if system_instruction is not None else {}), + **({"tools": _count_tokens_tools(generated_tools)} if generated_tools is not None else {}), + } + ) + + +def _tools(description: str) -> JsonValue: + return _JSON_LIST.validate_python( + [ + { + "name": "noop", + "description": description, + "input_schema": {"type": "object", "properties": {}}, + } + ] + ) + + +@pytest.mark.parametrize( + ("large_system", "large_tool"), + ((False, False), (True, False), (False, True), (True, True)), + ids=("baseline", "large-system", "large-tool", "large-system-and-tool"), +) +def test_messages_and_count_tokens_send_matching_gemini_context( + gateway: Gateway, large_system: bool, large_tool: bool +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + system: Final = ( + f"{uuid4().hex}\n" + "This is a synthetic context-budget diagnostic. Return exactly OK." + (_UNIT * 100 if large_system else "") + if large_system + else None + ) + tools: Final = _tools(_UNIT * 100) if large_tool else None + body: Final = _messages(model, system, tools) + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + requests: Final = peer.drain() + generated_body: Final = _json_body( + next(request for request in requests if request.target.endswith(":generateContent")) + ) + count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens"))) + assert count_body == _count_tokens_request_from_generate_content(generated_body) + + +def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + system: Final[JsonValue] = [ + {"type": "text", "text": "First instruction"}, + {"type": "text", "text": "Second instruction"}, + ] + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, system, None)) + assert response.status_code == 200, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + assert _json_body(requests[0]) == { + "contents": [{"role": "user", "parts": [{"text": _USER}]}], + "systemInstruction": {"parts": [{"text": "First instruction"}, {"text": "Second instruction"}]}, + } + + +def test_count_tokens_without_system_or_tools_sends_contents_only(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, None)) + assert response.status_code == 200, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + assert _json_body(requests[0]) == {"contents": [{"role": "user", "parts": [{"text": _USER}]}]} + + +_SDK_COUNT_ARGS: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "messages": [{"role": "user", "content": _USER}], + "system": "Follow the system instruction.", + "tools": [ + { + "name": "noop", + "description": "Return no output", + "input_schema": {"type": "object", "properties": {}}, + } + ], + } +) + + +def _sync_sdk_count_tokens(gateway: Gateway, model: str) -> int: + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + return client.messages.count_tokens(model=model, **_SDK_COUNT_ARGS).input_tokens # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping + + +async def _async_sdk_count_tokens(gateway: Gateway, model: str) -> int: + client: Final = anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + return (await client.messages.count_tokens(model=model, **_SDK_COUNT_ARGS)).input_tokens # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync-sdk", "async-sdk")) +def test_anthropic_sdk_count_tokens_forwards_system_and_tools(gateway: Gateway, asynchronous: bool) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + input_tokens: Final = ( + asyncio.run(_async_sdk_count_tokens(gateway, model)) + if asynchronous + else _sync_sdk_count_tokens(gateway, model) + ) + assert input_tokens == _COUNT + requests: Final = _count_requests(wire) + assert len(requests) == 1 + request_body: Final = _json_body(requests[0]) + assert request_body["systemInstruction"] == {"parts": [{"text": "Follow the system instruction."}]} + assert request_body["tools"] == [ + { + "functionDeclarations": [ + {"name": "noop", "description": "Return no output", "parameters": {"type": "object"}} + ] + } + ] + + +def test_google_count_tokens_forwards_system_and_tools_to_vertex(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + body: Final[Mapping[str, JsonValue]] = { + "contents": [ + {"role": "user", "parts": [{"text": _USER}]}, + {"role": "model", "parts": [{"functionCall": {"name": "noop", "args": {"query": "value"}}}]}, + { + "role": "user", + "parts": [{"functionResponse": {"name": "noop", "response": {"result": "ok"}}}], + }, + ], + "systemInstruction": {"parts": [{"text": "Follow this instruction."}]}, + "tools": [ + { + "functionDeclarations": [ + { + "name": "noop", + "description": "Return no output", + "parameters": {"type": "OBJECT", "properties": {}}, + } + ] + } + ], + } + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["totalTokens"] == _COUNT, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + request_body: Final = _json_body(requests[0]) + assert request_body["contents"] == body["contents"] + assert request_body["systemInstruction"] == body["systemInstruction"] + assert request_body["tools"] == [ + { + "functionDeclarations": [ + { + "name": "noop", + "description": "Return no output", + "parameters": {"type": "object"}, + } + ] + } + ] + + +def test_responses_input_tokens_maps_flat_function_tools_to_vertex(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + body: Final[Mapping[str, JsonValue]] = { + "model": model, + "instructions": "Follow the system instruction.", + "input": [ + { + "type": "function_call", + "call_id": "call_lookup", + "name": "lookup", + "arguments": '{"query":"value"}', + }, + { + "type": "function_call_output", + "call_id": "call_lookup", + "output": "result", + }, + ], + "tools": [ + { + "type": "function", + "name": "lookup", + "description": "Look up a value", + "parameters": {"type": "object", "properties": {}}, + } + ], + } + response: Final = gateway.request("POST", "/v1/responses/input_tokens", body) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + request_body: Final = _json_body(requests[0]) + assert request_body["tools"] == [ + { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Look up a value", + "parameters": {"type": "object"}, + } + ] + } + ] + + +def test_utils_token_counter_maps_openai_function_tools_to_vertex(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + tool_schema: Final[JsonValue] = {"type": "object", "properties": {}} + body: Final[Mapping[str, JsonValue]] = { + "model": model, + "messages": _MESSAGES, + "system": "Follow the system instruction.", + "tools": [ + { + "type": "function", + "function": { + "name": "noop", + "description": "Return no output", + "parameters": tool_schema, + }, + } + ], + } + response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "true"}) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["total_tokens"] == _COUNT, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + request_body: Final = _json_body(requests[0]) + assert request_body["systemInstruction"] == {"parts": [{"text": "Follow the system instruction."}]} + assert request_body["tools"] == [ + { + "functionDeclarations": [ + { + "name": "noop", + "description": "Return no output", + "parameters": {"type": "object"}, + } + ] + } + ] + + +def test_count_tokens_surfaces_peer_errors_and_rejects_non_text_system(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + peer.set_status(400) + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, None)) + assert response.status_code == 400, response.text + assert len(_count_requests(wire)) == 1 + + invalid_peer: Final = _GeminiPeer() + with wire_server(invalid_peer) as invalid_wire, gateway.scenario() as scenario: + invalid_model: Final = _deployment(gateway, scenario, invalid_wire.url) + invalid_response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + _messages( + invalid_model, + [{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + None, + ), + ) + assert invalid_response.status_code == 400, invalid_response.text + assert _count_requests(invalid_wire) == () + + +@pytest.mark.parametrize( + "tools", + ( + [{"name": "noop", "input_schema": {"type": "object", "properties": {"query": "not-a-schema"}}}], + [{"name": "noop", "input_schema": "not-a-schema"}], + ), + ids=("malformed-property", "string-schema"), +) +def test_count_tokens_rejects_unconvertible_tools_without_request(gateway: Gateway, tools: JsonValue) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, tools)) + assert response.status_code == 400, response.text + assert "tools" in response.text + assert _count_requests(wire) == () + + +def _count_wave(clients: Sequence[httpx.Client], body: Mapping[str, JsonValue]) -> tuple[tuple[int, JsonValue], ...]: + def count(client: httpx.Client) -> tuple[int, JsonValue]: + response: Final = client.post("/v1/messages/count_tokens", json=dict(body)) + payload: Final = _JSON_OBJECT.validate_json(response.content) + return response.status_code, payload.get("input_tokens") + + with ThreadPoolExecutor(max_workers=len(clients)) as executor: + return tuple(executor.map(count, clients)) + + +def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario, ExitStack() as stack: + model: Final = _deployment(gateway, scenario, wire.url) + clients: Final = tuple( + stack.enter_context( + httpx.Client( + base_url=str(gateway.client.base_url), + headers={"Authorization": f"Bearer {gateway.key}"}, + timeout=30, + trust_env=False, + ) + ) + for _ in range(3) + ) + body: Final = _messages(model, None, None) + assert _count_wave(clients, body) == ((200, _COUNT), (200, _COUNT), (200, _COUNT)) + peer.set_status(503) + assert _count_wave(clients, body) == ((503, None), (503, None), (503, None)) + peer.set_status(200) + assert _count_wave(clients, body) == ((200, _COUNT), (200, _COUNT), (200, _COUNT)) + requests: Final = _count_requests(wire) + assert len(requests) == 9 + assert ( + tuple(_json_body(request) for request in requests) + == ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9 + ) diff --git a/tests/test_litellm/proxy/google_endpoints/test_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_endpoints.py index 92001118e2c..17af52b5579 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_endpoints.py @@ -2,15 +2,85 @@ Test for google_endpoints/endpoints.py """ +from typing import Final + import pytest -import sys, os from dotenv import load_dotenv - - -from litellm.proxy.google_endpoints.endpoints import google_count_tokens -from litellm.types.llms.vertex_ai import TokenCountDetailsResponse +from pydantic import JsonValue from starlette.requests import Request +from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request, google_count_tokens +from litellm.types.llms.vertex_ai import TokenCountDetailsResponse + + +@pytest.mark.parametrize( + ("system_key", "schema_key"), + (("systemInstruction", "parameters"), ("system_instruction", "parametersJsonSchema")), +) +def test_google_count_tokens_translates_system_and_function_declarations( + system_key: str, schema_key: str +) -> None: + schema: Final = {"type": "object", "properties": {"query": {"type": "string"}}} + request_body: Final[dict[str, JsonValue]] = { + "contents": [{"role": "user", "parts": [{"text": "Look this up"}]}], + system_key: {"parts": [{"text": "Use the provided tool."}]}, + "tools": [ + { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Find a record", + schema_key: schema, + } + ] + } + ], + } + request: Final = _google_count_tokens_request("claude-sonnet", request_body) + + assert request.system == [{"type": "text", "text": "Use the provided tool."}] + assert request.tools == [ + { + "name": "lookup", + "description": "Find a record", + "input_schema": schema, + } + ] + assert request.messages == [{"role": "user", "content": "Look this up"}] + + +@pytest.mark.parametrize( + ("schema_key", "expected_schema"), + ( + ("parameters", {"type": "object", "properties": {"query": {"type": "string"}}}), + ("parametersJsonSchema", {"type": "OBJECT", "properties": {"query": {"type": "STRING"}}}), + ), +) +def test_google_count_tokens_normalizes_parameters_only( + schema_key: str, expected_schema: dict[str, JsonValue] +) -> None: + request: Final = _google_count_tokens_request( + "claude-sonnet", + { + "contents": [{"role": "user", "parts": [{"text": "Look this up"}]}], + "tools": [ + { + "functionDeclarations": [ + { + "name": "lookup", + schema_key: { + "type": "OBJECT", + "properties": {"query": {"type": "STRING"}}, + }, + } + ] + } + ], + }, + ) + + assert request.tools == [{"name": "lookup", "input_schema": expected_schema}] + load_dotenv() diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index a3dd8021f73..8a21027872a 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -93,8 +93,10 @@ async def test_bedrock_runtime_count_is_kept_when_it_answers() -> None: @pytest.mark.parametrize( ("model_to_use", "runtime"), ( + ("global.anthropic.claude-opus-4-8", (401, {"Message": "not authorized to perform: bedrock:CountTokens"})), ("global.anthropic.claude-opus-4-8", (403, {"Message": "not authorized to perform: bedrock:CountTokens"})), ("amazon.nova-pro-v1:0", (400, UNSUPPORTED)), + ("amazon.nova-pro-v1:0", (404, UNSUPPORTED)), ), ) async def test_other_bedrock_runtime_failures_are_not_retried_on_mantle( @@ -120,7 +122,7 @@ async def test_other_bedrock_runtime_failures_are_not_retried_on_mantle( @pytest.mark.asyncio async def test_mantle_failure_is_reported_with_its_status() -> None: bedrock: Final = _Bedrock( - runtime=(400, UNSUPPORTED), + runtime=(404, UNSUPPORTED), mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}), ) @@ -138,3 +140,37 @@ async def test_mantle_failure_is_reported_with_its_status() -> None: assert result.tokenizer_type == "bedrock_mantle_api" assert "does not exist" in (result.error_message or "") assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) + + +@pytest.mark.asyncio +async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> None: + bedrock: Final = _Bedrock(runtime=(404, UNSUPPORTED), mantle=(200, {"input_tokens": 2177})) + tools: Final = [ + { + "name": "lookup", + "description": "Find a record", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ] + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + system="You are a terse assistant.", + tools=tools, + ) + + assert result is not None + assert result.error is False + assert result.total_tokens == 2177 + assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) + mantle_body: Final = _JSON_BODY.validate_json(bedrock.requests[1].content) + assert mantle_body == { + "model": "anthropic.claude-opus-4-8", + "messages": MESSAGES, + "system": "You are a terse assistant.", + "tools": tools, + } diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 70946e10590..9a47d01ec46 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -1,8 +1,149 @@ +from typing import Final from unittest.mock import AsyncMock, patch +import httpx import pytest +from pydantic import JsonValue, TypeAdapter +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.common_utils import GeminiModelInfo, GoogleAIStudioTokenCounter +from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler +from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter as GeminiCountTokensHandler + +_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue]) + + +class _GeminiCountTokensPeer: + def __init__(self) -> None: + self.request_body: dict[str, JsonValue] | None = None + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.request_body = _JSON_BODY.validate_json(request.content) + return httpx.Response(200, json={"totalTokens": 11}, request=request) + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_forwards_system_and_tools() -> None: + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + deployment={ + "litellm_params": { + "api_key": "test-key", + "api_base": "https://count-tokens.test", + "system_instruction": {"parts": [{"text": "deployment-level"}]}, + } + }, + request_model="gemini/gemini-2.5-flash", + tools=[ + { + "name": "lookup", + "description": "Find a record", + "input_schema": {"type": "object", "properties": {}}, + } + ], + system="You are terse.", + ) + + assert result is not None + assert result.error is False + assert result.total_tokens == 11 + assert peer.request_body == { + "generateContentRequest": { + "model": "models/gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "systemInstruction": {"parts": [{"text": "You are terse."}]}, + "tools": [ + { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Find a record", + "parameters": {"type": "object"}, + } + ] + } + ], + } + } + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_rejects_image_system_without_request() -> None: + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + system=[ + {"type": "text", "text": "You are terse."}, + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}, + }, + ], + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert peer.request_body is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("system_instruction", "tools", "expected_body"), + ( + ( + {"parts": [{"text": "Follow the system rule."}]}, + [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "OBJECT"}}]}], + { + "generateContentRequest": { + "model": "models/gemini-4-argon", + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "systemInstruction": {"parts": [{"text": "Follow the system rule."}]}, + "tools": [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "OBJECT"}}]}], + } + }, + ), + ( + None, + [], + {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}, + ), + ), +) +async def test_gemini_count_tokens_sends_optional_request_fields( + system_instruction: dict[str, JsonValue] | None, + tools: list[dict[str, JsonValue]], + expected_body: dict[str, JsonValue], +) -> None: + peer: Final = _GeminiCountTokensPeer() + token_counter: Final = GeminiCountTokensHandler( + client=AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + ) + + result: Final = await token_counter.acount_tokens( + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + model="gemini-4-argon", + api_key="test-key", + api_base="https://count-tokens.test", + system_instruction=system_instruction, + tools=tools, + ) + + assert result == {"totalTokens": 11} + assert peer.request_body == expected_body class TestGeminiModelInfo: @@ -158,7 +299,10 @@ class TestGoogleAIStudioTokenCounter: # Verify the mock was called correctly mock_acount_tokens.assert_called_once_with( - model=model_to_use, contents=contents + model=model_to_use, + contents=contents, + system_instruction=None, + tools=None, ) def test_clean_contents_for_gemini_api_removes_id_field(self): diff --git a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py index c2efc1acdb9..98ddfa08018 100644 --- a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -1,3 +1,4 @@ +from typing import Final import pytest @@ -79,6 +80,54 @@ def test_transform_includes_tools(): assert result["tools"] == tools +_SCHEMA: Final[dict[str, object]] = {"type": "object", "properties": {}} # mutable-ok: matches the tools param type + + +@pytest.mark.parametrize( + ("tool", "expected"), + ( + ( + {"name": "lookup", "input_schema": _SCHEMA}, + {"type": "function", "name": "lookup", "description": "", "parameters": _SCHEMA, "strict": False}, + ), + ( + {"name": "lookup", "input_schema": _SCHEMA, "strict": True}, + {"type": "function", "name": "lookup", "description": "", "parameters": _SCHEMA, "strict": True}, + ), + ( + {"type": "custom", "name": "lookup", "description": "Find a record", "input_schema": _SCHEMA}, + { + "type": "function", + "name": "lookup", + "description": "Find a record", + "parameters": _SCHEMA, + "strict": False, + }, + ), + ), + ids=("anthropic", "anthropic-strict", "anthropic-custom-type"), +) +def test_transform_anthropic_tool_shapes(tool: dict[str, object], expected: dict[str, object]) -> None: + result: Final = OpenAICountTokensConfig().transform_request_to_count_tokens( + model="gpt-4o", input="Find a record", tools=[tool] + ) + + assert result["tools"] == [expected] + + +def test_transform_builtin_tool_passes_through(): + config: Final = OpenAICountTokensConfig() + tools: Final = [{"type": "web_search", "search_context_size": "low"}] + + result: Final = config.transform_request_to_count_tokens( + model="gpt-4o", + input="Find a record", + tools=tools, + ) + + assert result["tools"] == tools + + def test_transform_no_instructions_no_tools(): """Test that None values are not included.""" config = OpenAICountTokensConfig() @@ -102,9 +151,7 @@ def test_messages_to_responses_input_basic(): {"role": "user", "content": "How are you?"}, ] - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( - messages - ) + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) assert len(input_items) == 3 assert input_items[0] == {"role": "user", "content": "Hello"} @@ -120,9 +167,7 @@ def test_messages_to_responses_input_with_system(): {"role": "user", "content": "Hello"}, ] - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( - messages - ) + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) assert len(input_items) == 1 assert input_items[0] == {"role": "user", "content": "Hello"} @@ -136,9 +181,7 @@ def test_messages_to_responses_input_with_developer(): {"role": "user", "content": "Hello"}, ] - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( - messages - ) + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) assert len(input_items) == 1 assert instructions == "Be concise." @@ -151,9 +194,7 @@ def test_messages_to_responses_input_with_tool(): {"role": "tool", "content": "72°F", "tool_call_id": "call_123"}, ] - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( - messages - ) + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) assert len(input_items) == 2 assert input_items[1] == { @@ -424,10 +465,7 @@ def test_validate_request_missing_input(): def test_get_endpoint_default(): """Test default endpoint URL.""" config = OpenAICountTokensConfig() - assert ( - config.get_openai_count_tokens_endpoint() - == "https://api.openai.com/v1/responses/input_tokens" - ) + assert config.get_openai_count_tokens_endpoint() == "https://api.openai.com/v1/responses/input_tokens" def test_get_endpoint_custom_base(): diff --git a/tests/unit/llms/vertex_ai/count_tokens/__init__.py b/tests/unit/llms/vertex_ai/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py new file mode 100644 index 00000000000..1c05a814fa4 --- /dev/null +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -0,0 +1,196 @@ +from copy import deepcopy +from typing import Final + +import pytest +from pydantic import JsonValue + +from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + truncate_tool_name, +) +from litellm.llms.vertex_ai.count_tokens.context import ( + VertexGeminiCountContext, + VertexGeminiCountContextError, + vertex_gemini_count_context, +) +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + +_SCHEMA: Final[JsonValue] = { + "type": "object", + "properties": {}, + "additionalProperties": False, +} +_FUNCTION_DECLARATION: Final = { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Look up a value", + "parameters": {"type": "object"}, + } + ] +} + + +@pytest.mark.parametrize( + "tool", + ( + { + "name": "lookup", + "description": "Look up a value", + "input_schema": _SCHEMA, + }, + { + "type": "function", + "function": { + "name": "lookup", + "description": "Look up a value", + "parameters": _SCHEMA, + }, + }, + { + "type": "function", + "name": "lookup", + "description": "Look up a value", + "parameters": _SCHEMA, + }, + ), + ids=("anthropic", "openai-chat", "responses"), +) +def test_function_tool_shapes_map_to_vertex_declarations(tool: JsonValue) -> None: + context: Final = vertex_gemini_count_context(system=None, tools=[tool]) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools == (_FUNCTION_DECLARATION,) + assert _SCHEMA["additionalProperties"] is False + + +def test_long_anthropic_tool_name_is_truncated_for_vertex() -> None: + name: Final = "mcp__docs__" + "a" * 70 + context: Final = vertex_gemini_count_context( + system=None, + tools=[{"name": name, "description": "Look up a value", "input_schema": _SCHEMA}], + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools == ( + { + "functionDeclarations": [ + { + "name": truncate_tool_name(name), + "description": "Look up a value", + "parameters": {"type": "object"}, + } + ] + }, + ) + assert len(truncate_tool_name(name)) == 64 + + +def test_64_character_anthropic_tool_name_is_unchanged_for_vertex() -> None: + name: Final = "mcp__docs__" + "a" * 53 + context: Final = vertex_gemini_count_context( + system=None, + tools=[{"name": name, "input_schema": _SCHEMA}], + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools == ({"functionDeclarations": [{"name": name, "parameters": {"type": "object"}}]},) + + +def test_long_openai_chat_tool_name_is_unchanged_for_vertex() -> None: + name: Final = "mcp__docs__" + "a" * 70 + context: Final = vertex_gemini_count_context( + system=None, + tools=[{"type": "function", "function": {"name": name, "parameters": _SCHEMA}}], + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools == ({"functionDeclarations": [{"name": name, "parameters": {"type": "object"}}]},) + + +def test_unknown_tools_are_omitted() -> None: + context: Final = vertex_gemini_count_context( + system=None, + tools=[{"type": "web_search_preview"}, {"type": "computer_use"}], + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools is None + + +def test_text_block_system_instruction_keeps_one_part_per_block() -> None: + context: Final = vertex_gemini_count_context( + system=[{"type": "text", "text": "First instruction."}, {"type": "text", "text": "Second instruction."}], + tools=None, + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.system_instruction == {"parts": ({"text": "First instruction."}, {"text": "Second instruction."})} + + +def test_string_system_instruction_maps_to_one_part() -> None: + context: Final = vertex_gemini_count_context(system="Follow this instruction.", tools=None) + + assert isinstance(context, VertexGeminiCountContext) + assert context.system_instruction == {"parts": ({"text": "Follow this instruction."},)} + + +def test_non_text_system_block_returns_an_error_value() -> None: + context: Final = vertex_gemini_count_context( + system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + tools=None, + ) + + assert isinstance(context, VertexGeminiCountContextError) + assert "text blocks" in context.error_message + + +@pytest.mark.parametrize(("system", "tools"), ((None, None), ("", None), ([], []))) +def test_empty_context_is_omitted(system: JsonValue | None, tools: JsonValue | None) -> None: + context: Final = vertex_gemini_count_context(system=system, tools=tools) + + assert isinstance(context, VertexGeminiCountContext) + assert context.system_instruction is None + assert context.tools is None + + +@pytest.mark.parametrize( + "tool", + ( + {"description": "Look up a value", "input_schema": _SCHEMA}, + {"type": "bash_20250124", "name": "bash"}, + ), + ids=("unnamed", "anthropic-hosted"), +) +def test_anthropic_tools_count_the_same_declarations_inference_sends(tool: dict[str, JsonValue]) -> None: + translated_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai([deepcopy(tool)]) # pyright: ignore[reportArgumentType] # raw Anthropic tool dict + inference_tools: Final = VertexGeminiConfig()._map_function( # pyright: ignore[reportPrivateUsage] # the tool mapping inference uses + value=[dict(translated) for translated in translated_tools], + optional_params={}, + ) + + context: Final = vertex_gemini_count_context(system=None, tools=[tool]) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools is not None + assert [tool_value["functionDeclarations"] for tool_value in context.tools] == [ + inference_tool["function_declarations"] for inference_tool in inference_tools + ] + + +@pytest.mark.parametrize( + "tools", + ( + "lookup", + [1], + [{"name": "lookup", "input_schema": "not-a-schema"}], + [{"type": "function", "function": {"name": "lookup", "parameters": [1]}}], + [{"name": "lookup", "input_schema": {"type": "object", "properties": {"query": "not-a-schema"}}}], + ), + ids=("not-a-list", "non-object-tool", "string-schema", "list-parameters", "malformed-property"), +) +def test_invalid_tools_return_an_error_value(tools: JsonValue) -> None: + context: Final = vertex_gemini_count_context(system=None, tools=tools) + + assert isinstance(context, VertexGeminiCountContextError) + assert "tools" in context.error_message diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 86eb26a4c15..59cd53567ae 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1,11 +1,15 @@ -from unittest.mock import patch +from typing import Final +from unittest.mock import AsyncMock, patch +import httpx import pytest +from pydantic import JsonValue, TypeAdapter from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH - - +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler from litellm.llms.vertex_ai.common_utils import ( + VertexAITokenCounter, _get_vertex_url, convert_anyof_null_to_nullable, get_vertex_location_from_url, @@ -16,6 +20,23 @@ from litellm.llms.vertex_ai.common_utils import ( validate_vertex_location, vertex_request_labels_from_litellm_params, ) +from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler + + +@pytest.mark.asyncio +async def test_vertex_gemini_rejects_non_text_system_blocks_before_request() -> None: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-4-argon", + messages=[{"role": "user", "content": "Count this"}], + contents=None, + request_model="vertex_ai/gemini-4-argon", + system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "text blocks" in (result.error_message or "") @pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"]) @@ -149,15 +170,11 @@ async def test_get_supports_system_message(): from litellm.llms.vertex_ai.common_utils import get_supports_system_message # fine-tuned vertex gemini models will specifiy they are in the /gemini spec format - result = get_supports_system_message( - model="gemini/1234567890", custom_llm_provider="vertex_ai" - ) + result = get_supports_system_message(model="gemini/1234567890", custom_llm_provider="vertex_ai") assert result == True # non-fine-tuned vertex gemini models will not specifiy they are in the /gemini spec format - result = get_supports_system_message( - model="random-model-name", custom_llm_provider="vertex_ai" - ) + result = get_supports_system_message(model="random-model-name", custom_llm_provider="vertex_ai") assert result == False @@ -230,13 +247,9 @@ def test_build_vertex_schema(): "properties": { "tags": {"items": {"type": "string"}, "type": "array"}, "metadata": {"type": "object"}, - "callbacks": { - "anyOf": [{"items": {}, "type": "array"}, {}, {"type": "null"}] - }, + "callbacks": {"anyOf": [{"items": {}, "type": "array"}, {}, {"type": "null"}]}, "run_name": {"type": "string"}, - "max_concurrency": { - "anyOf": [{"type": "integer"}, {"type": "null"}] - }, + "max_concurrency": {"anyOf": [{"type": "integer"}, {"type": "null"}]}, "recursion_limit": {"type": "integer"}, "configurable": {"type": "object"}, "run_id": { @@ -280,9 +293,7 @@ def test_build_vertex_schema(): ] }, "run_name": {"type": "string"}, - "max_concurrency": { - "anyOf": [{"type": "integer", "nullable": True}] - }, + "max_concurrency": {"anyOf": [{"type": "integer", "nullable": True}]}, "recursion_limit": {"type": "integer"}, "configurable": {"type": "object"}, "run_id": {"anyOf": [{"type": "string", "nullable": True}]}, @@ -689,9 +700,7 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix): ({}, "us-east1", "us-east1"), ], ) -def test_get_vertex_region_global_only_model( - model_cost_entry, vertex_region, expected_region -): +def test_get_vertex_region_global_only_model(model_cost_entry, vertex_region, expected_region): """Test get_vertex_region resolves region from model_cost supported_regions""" import litellm from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -703,9 +712,7 @@ def test_get_vertex_region_global_only_model( {"vertex_ai/test-model": model_cost_entry}, clear=False, ): - result = vertex_base.get_vertex_region( - vertex_region=vertex_region, model="test-model" - ) + result = vertex_base.get_vertex_region(vertex_region=vertex_region, model="test-model") assert result == expected_region @@ -823,9 +830,7 @@ def test_convert_schema_types_type_array_conversion(): assert anyof_types[1]["type"] == "number" # 4. Other properties preserved - assert ( - input_schema["properties"]["studio"]["description"] == "The studio ID or name" - ) + assert input_schema["properties"]["studio"]["description"] == "The studio ID or name" assert input_schema["required"] == ["studio"] @@ -992,7 +997,9 @@ def test_construct_target_url_with_version_prefix(): ), ], ) -def test_construct_target_url_versionless_project_route_gets_api_version(requested_route: str, expected_url: str) -> None: +def test_construct_target_url_versionless_project_route_gets_api_version( + requested_route: str, expected_url: str +) -> None: from litellm.llms.vertex_ai.common_utils import construct_target_url target_url = construct_target_url( @@ -1125,10 +1132,7 @@ def test_fix_enum_types(): # 2. Non-string enums are removed assert "enum" not in input_schema["properties"]["maxLength"] assert "enum" not in input_schema["properties"]["enabled"] - assert ( - "enum" - not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"] - ) + assert "enum" not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"] # 3. anyOf with string type keeps enum, non-string removes it assert "enum" in input_schema["properties"]["anyOfField"]["anyOf"][0] @@ -1182,8 +1186,6 @@ def test_get_token_url(): assert "v1beta1" not in url assert "/v1/" in url - pass - @pytest.mark.asyncio async def test_vertex_ai_token_counter_routes_partner_models(): @@ -1461,6 +1463,49 @@ async def test_vertex_ai_token_counter_routes_gemini_models(): assert result.total_tokens == 50 +@pytest.mark.asyncio +async def test_vertex_gemini_count_tokens_uses_request_system_over_deployment_system_instruction() -> None: + class VertexCountTokensPeer: + def __init__(self) -> None: + self.request_body: dict[str, JsonValue] | None = None + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.request_body = TypeAdapter(dict[str, JsonValue]).validate_json(request.content) + return httpx.Response(200, json={"totalTokens": 37}, request=request) + + peer: Final = VertexCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with ( + patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), + patch.object( + vertex_count_tokens_handler.VertexAITokenCounter, + "_ensure_access_token_async", + new=AsyncMock(return_value=("fake-token", "test-project")), + ), + ): + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={ + "litellm_params": { + "vertex_project": "test-project", + "vertex_location": "us-central1", + "system_instruction": {"parts": [{"text": "deployment-level"}]}, + } + }, + request_model="vertex_ai/gemini-2.5-flash", + system="You are terse.", + ) + + assert result is not None + assert result.error is False + assert result.total_tokens == 37 + assert peer.request_body is not None + assert peer.request_body["systemInstruction"] == {"parts": [{"text": "You are terse."}]} + + @pytest.mark.asyncio async def test_vertex_ai_token_counter_converts_messages_to_contents_for_gemini(): """ @@ -1475,9 +1520,7 @@ async def test_vertex_ai_token_counter_converts_messages_to_contents_for_gemini( token_counter = VertexAITokenCounter() - with patch( - "litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens" - ) as mock_acount_tokens: + with patch("litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens") as mock_acount_tokens: mock_acount_tokens.return_value = { "totalTokens": 42, "tokenizer_used": "gemini", @@ -1519,9 +1562,7 @@ async def test_vertex_ai_token_counter_returns_none_when_api_omits_total_tokens( token_counter = VertexAITokenCounter() - with patch( - "litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens" - ) as mock_acount_tokens: + with patch("litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens") as mock_acount_tokens: mock_acount_tokens.return_value = {"tokenizer_used": "gemini"} result = await token_counter.count_tokens( @@ -1564,9 +1605,7 @@ async def test_vertex_ai_partner_model_detection(): # Test Minimax models assert VertexAIPartnerModels.is_vertex_partner_model("minimaxai/minimax-m2-maas") # Test Moonshot models - assert VertexAIPartnerModels.is_vertex_partner_model( - "moonshotai/kimi-k2-thinking-maas" - ) + assert VertexAIPartnerModels.is_vertex_partner_model("moonshotai/kimi-k2-thinking-maas") # Test Gemini models (should NOT be detected as partner model) assert not VertexAIPartnerModels.is_vertex_partner_model("gemini-1.5-pro") @@ -1597,9 +1636,7 @@ def test_vertex_ai_moonshot_uses_openai_handler(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.should_use_openai_handler( - "moonshotai/kimi-k2-thinking-maas" - ) + assert VertexAIPartnerModels.should_use_openai_handler("moonshotai/kimi-k2-thinking-maas") def test_vertex_ai_zai_uses_openai_handler(): @@ -1634,9 +1671,7 @@ def test_vertex_ai_gemma_maas_is_partner_model(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.is_vertex_partner_model( - "google/gemma-4-26b-a4b-it-maas" - ) + assert VertexAIPartnerModels.is_vertex_partner_model("google/gemma-4-26b-a4b-it-maas") def test_vertex_ai_gemma_maas_uses_openai_handler(): @@ -1647,9 +1682,7 @@ def test_vertex_ai_gemma_maas_uses_openai_handler(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.should_use_openai_handler( - "google/gemma-4-26b-a4b-it-maas" - ) + assert VertexAIPartnerModels.should_use_openai_handler("google/gemma-4-26b-a4b-it-maas") def test_vertex_ai_gemma_maas_routes_to_partner_models(): @@ -1851,12 +1884,8 @@ def test_pop_vertex_request_labels_prefers_explicit_labels_then_metadata(): def test_pop_vertex_request_labels_uses_litellm_metadata_when_metadata_absent(): optional: dict = {} - litellm_params = { - "litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}} - } - assert pop_vertex_request_labels(optional, litellm_params) == { - "team": "from_litellm_meta" - } + litellm_params = {"litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}}} + assert pop_vertex_request_labels(optional, litellm_params) == {"team": "from_litellm_meta"} def test_vertex_text_embedding_request_includes_labels_from_metadata(): @@ -1866,9 +1895,7 @@ def test_vertex_text_embedding_request_includes_labels_from_metadata(): input="hi", optional_params={}, model="text-embedding-004", - litellm_params={ - "metadata": {"requester_metadata": {"project_id": "cost-center-1"}} - }, + litellm_params={"metadata": {"requester_metadata": {"project_id": "cost-center-1"}}}, ) assert req.get("labels") == {"project_id": "cost-center-1"} From c99e4b86db2453d3199785a25c40964417562281 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:23:45 -0700 Subject: [PATCH 13/36] ci(cost map): classify token_count_framing and token_count_fallback_model in the schema generator (#46784) Backport of #46784 to stable/1.104.x. Cherry-picked from c6a4437e980e1a498066399f69666a14aa618d05 (main). --- ci_cd/generate_model_prices_schema.py | 9 +++++++ tests/unit/test_model_prices_schema.py | 33 ++++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 8eec07dadda..e429aae36b9 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -189,6 +189,10 @@ INTEGER_KEYS: dict[str, JsonSchema] = { **NONNEG_INTEGER, "description": "Smallest prefix the provider will actually cache; absent means the provider default applies.", }, + "token_count_framing": { + **NONNEG_INTEGER, + "description": "Tokens added once per request to a local count estimate, for the provider's chat framing.", + }, "tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."}, "rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."}, } @@ -283,6 +287,11 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]: "enum": ["none", "minimal", "low", "medium", "high", "xhigh"], }, "comment": STRING, + "token_count_fallback_model": { + "type": "string", + "minLength": 1, + "description": "Model that counts tokens for this one when the provider cannot count it.", + }, "audio_transcription_config": STRING, "vertex_ai_audio_api": { "type": "string", diff --git a/tests/unit/test_model_prices_schema.py b/tests/unit/test_model_prices_schema.py index a05c345b5cd..ccc379ffe4f 100644 --- a/tests/unit/test_model_prices_schema.py +++ b/tests/unit/test_model_prices_schema.py @@ -176,6 +176,39 @@ def test_generated_off_peak_schema_rejects_malformed_blocks(block: dict): assert not validator.is_valid({"some-model": {**OFF_PEAK_ENTRY, "off_peak_pricing": block}}) +TOKEN_COUNT_ENTRY: Final = MappingProxyType( + { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "token_count_framing": 78, + "token_count_fallback_model": "us.anthropic.claude-opus-5-5", + } +) + + +@pytest.mark.parametrize("framing", [0, 78]) +def test_generator_classifies_token_count_settings(framing: int): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) + assert build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, "token_count_framing": framing}}) + + +@pytest.mark.parametrize( + "setting", + [ + {"token_count_framing": -1}, + {"token_count_framing": "78"}, + {"token_count_framing": 1.5}, + {"token_count_fallback_model": ""}, + {"token_count_fallback_model": 5}, + ], +) +def test_generated_token_count_schema_rejects_malformed_settings(setting: dict): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) + assert not build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}}) + + def find_duplicate_keys(path: Path) -> list[str]: duplicates: list[str] = [] From 531b66719a57ed2e98a01aba6eab093f8f70f83c Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:32:21 -0700 Subject: [PATCH 14/36] fix(count_tokens): count Anthropic-format requests for OpenAI the way /v1/messages inference converts them (#46759) Backport of #46759 to stable/1.104.x. Cherry-picked from d25b58f744bc151b8dec89a11bd6158a4f1a67e9 (main). Co-authored-by: Mrinal Chanshetty --- .../anthropic/count_tokens/token_counter.py | 10 +- .../messages/utils.py | 13 + .../responses_adapters/transformation.py | 17 +- .../anthropic/count_tokens/token_counter.py | 6 +- litellm/llms/base_llm/base_utils.py | 6 +- .../count_tokens/bedrock_token_counter.py | 4 +- litellm/llms/gemini/common_utils.py | 5 +- .../openai/responses/count_tokens/handler.py | 7 +- .../responses/count_tokens/token_counter.py | 78 +++- .../responses/count_tokens/transformation.py | 88 +++- litellm/llms/vertex_ai/common_utils.py | 6 +- .../proxy/anthropic_endpoints/endpoints.py | 17 +- litellm/proxy/proxy_server.py | 16 +- litellm/types/utils.py | 4 + .../test_openai_count_tokens_tools_wire.py | 192 ++++++++- .../anthropic_endpoints/test_endpoints.py | 8 +- .../count_tokens/test_token_counter.py | 390 ++++++++++++++++++ ...test_openai_count_tokens_transformation.py | 108 +++++ tests/unit/proxy/test_proxy_token_counter.py | 159 +++---- tests/unit/test_count_tokens_public_api.py | 2 +- 20 files changed, 990 insertions(+), 146 deletions(-) create mode 100644 tests/unit/llms/openai/responses/count_tokens/test_token_counter.py diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 8e8d10c961b..1ca882071cb 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -5,10 +5,12 @@ Anthropic Token Counter implementation using the CountTokens API. import os from typing import Any, Final +from pydantic import JsonValue + from litellm._logging import verbose_logger from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler() @@ -30,8 +32,10 @@ class AnthropicTokenCounter(BaseTokenCounter): contents: list[dict[str, Any]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, + tools: list[dict[str, JsonValue]] | None = None, + system: JsonValue = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using Anthropic's CountTokens API. diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 89105c00428..f36d3a48726 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -113,6 +113,19 @@ def anthropic_system_to_openai_message(system: object) -> ChatCompletionSystemMe return ChatCompletionSystemMessage(role="system", content=system) +def normalize_native_messages(messages: Sequence[object]) -> list[dict[str, object]]: + """Apply the block rewrites every /v1/messages request gets before it reaches a provider.""" + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + sanitize_tool_use_ids_in_anthropic_messages, + strip_empty_content_blocks_from_anthropic_messages, + ) + + return flatten_unencrypted_web_search_results_in_anthropic_messages( + sanitize_tool_use_ids_in_anthropic_messages(strip_empty_content_blocks_from_anthropic_messages(messages)) + ) + + @lru_cache(maxsize=1) def _anthropic_messages_optional_param_keys() -> frozenset[str]: """ diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 6a31173a9c6..ba8bf25f6d3 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -500,6 +500,19 @@ class LiteLLMAnthropicToResponsesAPIAdapter: result["summary"] = "detailed" return result + @staticmethod + def system_blocks_to_instructions(system: object) -> str: + if not isinstance(system, list): + return "" + return "\n".join( + text + for block in system + if isinstance(block, dict) + and block.get("type") == "text" + and isinstance(text := block.get("text"), str) + and text + ) + def translate_request( self, anthropic_request: AnthropicMessagesRequest, @@ -549,9 +562,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(system, str): responses_kwargs["instructions"] = system elif isinstance(system, list): - responses_kwargs["instructions"] = "\n".join( - filter(None, (b.get("text", "") for b in system if isinstance(b, dict) and b.get("type") == "text")) - ) + responses_kwargs["instructions"] = self.system_blocks_to_instructions(system) # max_tokens -> max_output_tokens max_tokens: Final = anthropic_request.get("max_tokens") diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index c31d2c427bb..c02785f2051 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -5,12 +5,14 @@ Azure AI Anthropic Token Counter implementation using the CountTokens API. import os from typing import Any, Final +from pydantic import JsonValue + from litellm._logging import verbose_logger from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler() @@ -34,6 +36,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using Azure AI Anthropic's CountTokens API. diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index a446a721307..80b9c466dfb 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -8,11 +8,11 @@ from abc import ABC, abstractmethod from typing import Any, Final from openai.lib import _parsing, _pydantic -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue from litellm._logging import verbose_logger from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk -from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountResponse +from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountMessagesFormat, TokenCountResponse class BaseTokenCounter(ABC): @@ -26,6 +26,8 @@ class BaseTokenCounter(ABC): request_model: str = "", tools: list[dict[str, Any]] | None = None, system: Any | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: pass diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 3084b0fde07..7ff2fddbb57 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -16,7 +16,7 @@ from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountResponse +from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" @@ -107,6 +107,8 @@ class BedrockTokenCounter(BaseTokenCounter): request_model: str = "", tools: Sequence[Mapping[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: if not messages: return None diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 2c3b163496e..bd155c00c8f 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -6,6 +6,7 @@ from collections.abc import Mapping, Sequence from typing import Any, Final import httpx +from pydantic import JsonValue import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -13,7 +14,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = { "1:1": 1 / 1, @@ -491,6 +492,8 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: import copy diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 66782a83a19..0f4e4655575 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -5,6 +5,7 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint. """ import json +from collections.abc import Mapping, Sequence from typing import Any, Final import httpx @@ -30,9 +31,10 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): api_key: str, api_base: str | None = None, timeout: float | httpx.Timeout | None = None, - tools: list[dict[str, Any]] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, - ) -> dict[str, Any]: + text: Mapping[str, object] | None = None, + ) -> dict[str, object]: """ Handle a token counting request to OpenAI's Responses API. @@ -52,6 +54,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): input=input, tools=tools, instructions=instructions, + text=text, ) endpoint_url: Final = self.get_openai_count_tokens_endpoint(api_base) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index c05d943b7cf..59c812a0a32 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -3,9 +3,16 @@ OpenAI Token Counter implementation using the Responses API /input_tokens endpoi """ import os +from collections.abc import Mapping, Sequence from typing import Any, Final +from pydantic import JsonValue +from typing_extensions import assert_never + from litellm._logging import verbose_logger +from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.count_tokens.handler import ( @@ -13,8 +20,9 @@ from litellm.llms.openai.responses.count_tokens.handler import ( ) from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, + ResponsesCountRequestParts, ) -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests openai_count_tokens_handler: Final = OpenAICountTokensHandler() @@ -29,6 +37,37 @@ class OpenAITokenCounter(BaseTokenCounter): ) -> bool: return custom_llm_provider == LlmProviders.OPENAI.value + @staticmethod + def _system_instructions(system: object) -> str | None: + if system is None or isinstance(system, str): + return system + return LiteLLMAnthropicToResponsesAPIAdapter.system_blocks_to_instructions(system) or None + + @staticmethod + def _responses_request_parts( + model_to_use: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + system: object, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue, + ) -> ResponsesCountRequestParts: + match messages_format: + case "anthropic": + return OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model=model_to_use, messages=messages, system=system, tools=tools, output_format=output_format + ) + case "openai": + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( + [dict(message) for message in messages] + ) + fallback_instructions: Final = OpenAITokenCounter._system_instructions(system) + return OpenAICountTokensConfig.count_request_parts( + input_items, instructions if instructions is not None else fallback_instructions, tools + ) + case _: + assert_never(messages_format) + async def count_tokens( self, model_to_use: str, @@ -38,6 +77,8 @@ class OpenAITokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using OpenAI's Responses API /input_tokens endpoint. @@ -45,6 +86,22 @@ class OpenAITokenCounter(BaseTokenCounter): if not messages: return None + malformed_error: Final = ( + OpenAICountTokensConfig.malformed_anthropic_messages_error(messages) + if messages_format == "anthropic" + else None + ) + if malformed_error is not None: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=malformed_error, + status_code=400, + ) + deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) @@ -59,12 +116,14 @@ class OpenAITokenCounter(BaseTokenCounter): api_base: Final = litellm_params.get("api_base") - # Convert chat messages to Responses API input format - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) - - # Use system param if instructions not extracted from messages - if instructions is None and system is not None: - instructions = system if isinstance(system, str) else str(system) + input_items, instructions, counted_tools, text = self._responses_request_parts( + model_to_use=model_to_use, + messages=messages, + tools=tools, + system=system, + messages_format=messages_format, + output_format=output_format, + ) # If no input items were produced (e.g., system-only messages), fall back to local counting if not input_items: @@ -73,11 +132,12 @@ class OpenAITokenCounter(BaseTokenCounter): try: result: Final = await openai_count_tokens_handler.handle_count_tokens_request( model=model_to_use, - input=input_items if input_items is not None else [], + input=input_items, api_key=api_key, api_base=api_base, - tools=tools, + tools=counted_tools, instructions=instructions, + text=text, ) if result is not None: diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index fa46fdc69ef..3ea5cf5c4b7 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -5,10 +5,18 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp """ from collections.abc import Mapping, Sequence -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypeAlias +from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict +from litellm.llms.anthropic.experimental_pass_through.messages.utils import normalize_native_messages +from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema + class ResponsesInputTextPart(TypedDict): type: ReadOnly[Literal["input_text"]] @@ -31,6 +39,30 @@ ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart | Response ResponsesContentRole = Literal["user", "assistant"] +ResponsesCountRequestParts: TypeAlias = tuple[ + tuple[Mapping[str, object], ...], + str | None, + tuple[Mapping[str, object], ...] | None, + Mapping[str, object] | None, +] +_RESPONSES_COUNT_REQUEST_PARTS: Final = TypeAdapter(ResponsesCountRequestParts) +_TRANSLATED_REQUEST: Final = TypeAdapter(Mapping[str, object]) + + +class _AnthropicCountMessage(TypedDict): + role: ReadOnly[str] + + +_ANTHROPIC_COUNT_MESSAGES: Final = TypeAdapter(tuple[_AnthropicCountMessage, ...]) +_ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema) + + +def _anthropic_output_schema(output_format: JsonValue) -> AnthropicOutputSchema | None: + try: + return _ANTHROPIC_OUTPUT_SCHEMA.validate_python(output_format) + except ValidationError: + return None + def _chat_image_block_to_responses_part(image_url: object) -> ResponsesInputImagePart | None: url: Final = image_url.get("url") if isinstance(image_url, Mapping) else image_url @@ -118,8 +150,9 @@ class OpenAICountTokensConfig: self, model: str, input: str | Sequence[object], - tools: list[dict[str, Any]] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, + text: Mapping[str, object] | None = None, ) -> dict[str, object]: """ Transform request to OpenAI Responses API token counting format. @@ -135,7 +168,10 @@ class OpenAICountTokensConfig: request["instructions"] = instructions if tools is not None: - request["tools"] = self._transform_tools_for_responses_api(tools) + request["tools"] = self._transform_tools_for_responses_api([dict(tool) for tool in tools]) + + if text is not None: + request["text"] = text return request @@ -194,7 +230,7 @@ class OpenAICountTokensConfig: @staticmethod def messages_to_responses_input( messages: list[dict[str, Any]], - ) -> tuple: + ) -> tuple[list[dict[str, object]], str | None]: """ Convert standard chat messages format to OpenAI Responses API input format. @@ -257,3 +293,47 @@ class OpenAICountTokensConfig: instructions: Final = "\n".join(instructions_parts) if instructions_parts else None return input_items, instructions + + @staticmethod + def malformed_anthropic_messages_error(messages: Sequence[Mapping[str, object]]) -> str | None: + try: + _ANTHROPIC_COUNT_MESSAGES.validate_python(messages) + except ValidationError as e: + return "; ".join( + f"messages.{'.'.join(str(part) for part in error['loc'])}: {error['msg']}" for error in e.errors() + ) + return None + + @staticmethod + def anthropic_messages_to_responses_input( + model: str, + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + output_format: JsonValue = None, + ) -> ResponsesCountRequestParts: + """Convert with the adapter /v1/messages inference uses, so the count sees the same input it bills.""" + anthropic_request: Final = AnthropicMessagesRequest( + model=model, + messages=normalize_native_messages(messages), + system=system if isinstance(system, (str, list)) else None, + tools=[dict(tool) for tool in tools] if tools else None, + output_format=_anthropic_output_schema(output_format), + ) + translated: Final = _TRANSLATED_REQUEST.validate_python( + LiteLLMAnthropicToResponsesAPIAdapter().translate_request( + anthropic_request, include_encrypted_reasoning=False + ) + ) + return OpenAICountTokensConfig.count_request_parts( + ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(translated.get("input")), # pyright: ignore[reportPrivateUsage] # the input restore inference runs + translated.get("instructions"), + translated.get("tools"), + translated.get("text"), + ) + + @staticmethod + def count_request_parts( + input_items: object, instructions: object, tools: object, text: object = None + ) -> ResponsesCountRequestParts: + return _RESPONSES_COUNT_REQUEST_PARTS.validate_python((input_items, instructions, tools, text)) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 8643ae31e29..7d4b87d0228 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -7,7 +7,7 @@ from types import MappingProxyType from typing import Any, Final, Literal, cast, get_type_hints import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -22,7 +22,7 @@ from litellm.types.llms.vertex_ai import ( PartType, Schema, ) -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages VERTEX_SELF_DEPLOYED_ENDPOINT_UNSUPPORTED_PARAMS: Final = frozenset( @@ -1237,6 +1237,8 @@ class VertexAITokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: import copy diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index d9558b86e95..f166fe87b4a 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -2,10 +2,12 @@ Unified /v1/messages endpoint - (Anthropic Spec) """ +from collections.abc import Mapping from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse +from pydantic import JsonValue import litellm from litellm.anthropic_interface.exceptions import ( @@ -286,6 +288,14 @@ async def anthropic_response( ) +def _requested_output_format(data: Mapping[str, JsonValue]) -> JsonValue: + output_format: Final = data.get("output_format") + output_config: Final = data.get("output_config") + if isinstance(output_format, dict) or not isinstance(output_config, dict): + return output_format + return output_config.get("format") + + @router.post( "/v1/messages/count_tokens", tags=["[beta] Anthropic Messages Token Counting"], @@ -315,7 +325,7 @@ async def count_tokens( Returns: {"input_tokens": } """ - from litellm.proxy.proxy_server import token_counter as internal_token_counter + from litellm.proxy.proxy_server import count_request_tokens litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: @@ -342,10 +352,11 @@ async def count_tokens( system=data.get("system"), ) - # Call the internal token counter function with direct request flag set to False - token_response: Final = await internal_token_counter( + token_response: Final = await count_request_tokens( request=token_request, call_endpoint=True, + messages_format="anthropic", + output_format=_requested_output_format(data), ) _token_response_dict: dict = {} if isinstance(token_response, TokenCountResponse): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 76e01a874a2..f5810a96a9c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -162,6 +162,7 @@ from litellm.types.utils import ( ModelResponseStream, StreamingChoices, TextCompletionResponse, + TokenCountMessagesFormat, TokenCountResponse, echoed_cost_map_pricing_fields, is_server_derived_pricing_key, @@ -13619,6 +13620,8 @@ async def _try_provider_token_count( request_model: str, tools: list | None = None, system: str | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> Optional["TokenCountResponse"]: """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): @@ -13632,6 +13635,8 @@ async def _try_provider_token_count( request_model=request_model, tools=tools, system=system, + messages_format=messages_format, + output_format=output_format, ) except httpx.HTTPStatusError as e: error_message: Final = getattr(e, "message", None) or str(e) @@ -13674,8 +13679,15 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) Returns: TokenCountResponse """ - global llm_router + return await count_request_tokens(request=request, call_endpoint=call_endpoint) + +async def count_request_tokens( + request: TokenCountRequest, + call_endpoint: bool = False, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, +) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages contents: Final = request.contents @@ -13734,6 +13746,8 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) request_model=request.model, tools=tools, system=system, + messages_format=messages_format, + output_format=output_format, ) if result is not None: return result diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 9c676ef39a1..d3f856697e3 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -9,6 +9,7 @@ from typing import ( Any, Final, Literal, + TypeAlias, get_args, ) @@ -4217,6 +4218,9 @@ class LiteLLMLoggingBaseClass: pass +TokenCountMessagesFormat: TypeAlias = Literal["openai", "anthropic"] + + class TokenCountResponse(LiteLLMPydanticObjectBase): total_tokens: int request_model: str diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index bfae84e333d..3ab5826006f 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -1,5 +1,5 @@ import json -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from typing import Final import pytest @@ -36,6 +36,146 @@ def _count_response(_request: Request) -> Reply: return Reply(body=json.dumps({"input_tokens": _COUNT}).encode()) +_PDF_BASE64: Final = "JVBERi0xLjQK" +_PNG_BASE64: Final = "aW1hZ2U=" +_LOOKUP_TOOL: Final[dict[str, JsonValue]] = { + "name": "lookup", + "description": "Find a record", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, +} +_RESPONSES_REPLY: Final[Mapping[str, JsonValue]] = { + "id": "resp_count_parity", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_count_parity", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "OK", "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}, +} +_WRAPPED_REASONING: Final = "litellm_enc:bW9kZWxfaWQ6ZGVwbG95bWVudC0x;gAAAAoriginal" +_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}} +_ANTHROPIC_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = { + "structured-output": { + "output_config": {"format": {"type": "json_schema", "schema": _ANSWER_SCHEMA}}, + "messages": [{"role": "user", "content": "Answer in JSON."}], + }, + "affinity-wrapped-reasoning": { + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{_WRAPPED_REASONING}"}, + {"type": "text", "text": "OK"}, + ], + }, + {"role": "user", "content": "again"}, + ], + }, + "base64-pdf": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the attached contract."}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": _PDF_BASE64}, + }, + ], + } + ], + }, + "base64-image": { + "system": [{"type": "text", "text": "Describe images briefly."}], + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": _PNG_BASE64}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + }, + "replayed-web-search": { + "messages": [ + {"role": "user", "content": "What is the newest LiteLLM release?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "newest LiteLLM release"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": "LiteLLM ships a weekly release candidate and a stable line.", + } + ], + }, + {"type": "text", "text": "The newest release is in the release notes."}, + ], + }, + {"role": "user", "content": "Thanks."}, + ], + }, + "tool-turns": { + "system": "Use the lookup tool.", + "tools": [_LOOKUP_TOOL], + "messages": [ + {"role": "user", "content": "Find record 14."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Looking it up."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "14"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}], + }, + ], + }, +} + + +def _responses_peer(request: Request) -> Reply: + if request.target == "/v1/responses/input_tokens": + return _count_response(request) + if request.target == "/v1/responses": + return Reply(body=json.dumps(_RESPONSES_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": request.target}).encode()) + + +def _body_for(wire_requests: tuple[Request, ...], target: str) -> dict[str, JsonValue]: + matching: Final = tuple(request for request in wire_requests if request.target == target) + assert len(matching) == 1, [request.target for request in wire_requests] + return _JSON_OBJECT.validate_json(matching[0].body) + + +def _counted_fields(body: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + return {key: body[key] for key in ("input", "instructions", "tools", "text") if body.get(key) is not None} + + def _count_body(wire: Wire) -> dict[str, JsonValue]: requests: Final = wire.drain() assert len(requests) == 1, requests @@ -159,3 +299,53 @@ def test_utils_token_counter_keeps_chat_tool_conversion(counting_proxy: OwnedPro "parameters": _TOOL_SCHEMA, } ] + + +@pytest.mark.parametrize("case", sorted(_ANTHROPIC_BODIES)) +def test_messages_count_tokens_sends_the_input_v1_messages_inference_sends( + counting_proxy: OwnedProxy, case: str +) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + body: Final = {"model": model, **_ANTHROPIC_BODIES[case]} + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT + requests: Final = wire.drain() + inference_body: Final = _body_for(requests, "/v1/responses") + count_body: Final = _body_for(requests, "/v1/responses/input_tokens") + assert _counted_fields(count_body) == _counted_fields(inference_body) + assert count_body["model"] == inference_body["model"] + + +def test_messages_count_tokens_without_a_role_answers_400_before_calling_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": [{"role": "user", "content": "hello"}, {"content": "no role"}]}, + ) + assert response.status_code == 400, response.text + assert "messages.1.role: Field required" in response.text + assert wire.drain() == () + + +def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + { + "systemInstruction": {"parts": [{"text": "Be brief."}, {"text": "Answer in English."}]}, + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + }, + ) + assert response.status_code == 200, response.text + assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English." diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index 801f61aa498..9ecf681e039 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -284,7 +284,9 @@ class TestHttpExceptionDictDetail: "base_process_llm_request", new=AsyncMock(side_effect=exc), ), - patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + patch.object( + proxy_server, "proxy_logging_obj" + ) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam ): mock_logging.post_call_failure_hook = AsyncMock() response = await ep.anthropic_response( @@ -440,7 +442,9 @@ class TestErrorLogCarriesCallId: "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]}), ), - patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block + patch.object( + proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down")) + ), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(HTTPException) as raised, ): diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py new file mode 100644 index 00000000000..0f8d2aaba54 --- /dev/null +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -0,0 +1,390 @@ +import json +import operator +from typing import Final + +import httpx +import pytest +import respx +from pydantic import JsonValue, ValidationError + +import litellm +from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse + +COUNT_URL: Final = "https://api.openai.com/v1/responses/input_tokens" + + +async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + respx.post(COUNT_URL).mock(return_value=response) + return await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "expected_total"), + [ + ({"input_tokens": 15}, 15), + ({"input_tokens": "15"}, 15), + ({"input_tokens": 15.0}, 15), + ({"input_tokens": True}, 1), + ({"input_tokens": -3}, -3), + ({}, 0), + ({"object": "response.input_tokens", "error": {"message": "429 rate limit"}}, 0), + ({"input_tokens": 7, "object": "response.input_tokens", "nested": {"a": [1, None]}}, 7), + ], +) +async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload( + monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": expected_total, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "original_response": payload, + "error": False, + "error_message": None, + "status_code": None, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rejected_count", + [ + None, + 15.5, + "Request timed out", + "429 Too Many Requests", + [15], + [*range(403), "Request timed out"], + {"value": 429}, + ], +) +async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message( + monkeypatch: pytest.MonkeyPatch, rejected_count: object +) -> None: + with pytest.raises(ValidationError) as rejection: + TokenCountResponse.model_validate( + { + "total_tokens": rejected_count, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + } + ) + + result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count})) + + assert result is not None + assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)] + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "original_response": None, + "error": True, + "error_message": str(rejection.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False], +) +async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup( + monkeypatch: pytest.MonkeyPatch, payload: object +) -> None: + with pytest.raises(AttributeError) as failed_lookup: + operator.attrgetter("get")(payload) + + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "original_response": None, + "error": True, + "error_message": str(failed_lookup.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None: + assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("upstream_status", "upstream_text"), + [(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")], +) +async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text( + monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "original_response": None, + "error": True, + "error_message": upstream_text, + "status_code": upstream_status, + } + ) + + +async def _sent_count_body( + monkeypatch: pytest.MonkeyPatch, + messages: list[dict[str, object]], + system: object, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue = None, +) -> object: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1})) + await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=messages, + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + system=system, + messages_format=messages_format, + output_format=output_format, + ) + return json.loads(route.calls.last.request.content) + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_the_document_image_and_tool_turns_to_the_count_endpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + messages: Final[list[dict[str, object]]] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize."}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBE"}}, + {"type": "image", "source": {"type": "url", "url": "https://example.com/page.png"}}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "found"}]}, + ] + + sent: Final = await _sent_count_body( + monkeypatch, messages, [{"type": "text", "text": "Be brief."}], messages_format="anthropic" + ) + + assert sent == { + "model": "gpt-4o", + "input": [ + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize."}, + {"type": "input_file", "filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBE"}, + {"type": "input_image", "image_url": "https://example.com/page.png"}, + ], + }, + {"type": "function_call", "call_id": "toolu_01", "name": "lookup", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "toolu_01", "output": "found"}, + ], + "instructions": "Be brief.", + } + + +@pytest.mark.asyncio +async def test_openai_format_keeps_the_chat_conversion(monkeypatch: pytest.MonkeyPatch) -> None: + messages: Final[list[dict[str, object]]] = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": [{"type": "text", "text": "Summarize."}]}, + {"role": "tool", "tool_call_id": "call_1", "content": "found"}, + ] + + sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="openai") + + assert sent == { + "model": "gpt-4o", + "input": [ + {"role": "user", "content": "Summarize."}, + {"type": "function_call_output", "call_id": "call_1", "output": "found"}, + ], + "instructions": "Be brief.", + } + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_affinity_wrapped_reasoning_unwrapped_as_inference_does( + monkeypatch: pytest.MonkeyPatch, +) -> None: + wrapped: Final = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAoriginal", "deployment-1") # pyright: ignore[reportPrivateUsage] # the wrap inference applies + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"}, + {"type": "text", "text": "OK"}, + ], + }, + {"role": "user", "content": "again"}, + ] + + sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic") + + assert sent == { + "model": "gpt-4o", + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + {"type": "reasoning", "summary": [], "encrypted_content": "gAAAAoriginal"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "OK"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "again"}]}, + ], + } + + +@pytest.mark.asyncio +async def test_anthropic_format_counts_the_structured_output_schema(monkeypatch: pytest.MonkeyPatch) -> None: + schema: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}} + + sent: Final = await _sent_count_body( + monkeypatch, + [{"role": "user", "content": "Answer in JSON."}], + system=None, + messages_format="anthropic", + output_format={"type": "json_schema", "schema": schema}, + ) + + assert sent == { + "model": "gpt-4o", + "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Answer in JSON."}]}], + "text": {"format": {"type": "json_schema", "name": "structured_output", "schema": schema, "strict": False}}, + } + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_a_replayed_web_search_result_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + {"role": "user", "content": "Thanks."}, + ] + + sent: Final = json.dumps(await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic")) + + assert snippet in sent + assert "https://docs.litellm.ai/release_notes" in sent + + +@pytest.mark.asyncio +async def test_anthropic_format_message_without_a_role_answers_400_without_calling_openai( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1})) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}, {"content": "no role"}], + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + messages_format="anthropic", + ) + + assert route.called is False + assert result is not None + assert (result.error, result.status_code, result.error_message) == (True, 400, "messages.1.role: Field required") + + +@pytest.mark.asyncio +async def test_openai_format_sends_a_list_system_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None: + sent: Final = await _sent_count_body( + monkeypatch, + [{"role": "user", "content": "hi"}], + system=[{"type": "text", "text": "Be brief."}, {"type": "text", "text": "Answer in English."}], + messages_format="openai", + ) + + assert sent == { + "model": "gpt-4o", + "input": [{"role": "user", "content": "hi"}], + "instructions": "Be brief.\nAnswer in English.", + } + + +@pytest.mark.asyncio +async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with pytest.raises(json.JSONDecodeError) as undecodable: + json.loads("not json") + + result: Final = await _count(monkeypatch, httpx.Response(200, text="not json")) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "original_response": None, + "error": True, + "error_message": f"CountTokens processing error: {undecodable.value}", + "status_code": 500, + } + ) diff --git a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py index 98ddfa08018..ddcbfc61822 100644 --- a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -2,6 +2,9 @@ from typing import Final import pytest +from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ) @@ -484,3 +487,108 @@ def test_get_required_headers(): assert headers["Authorization"] == "Bearer sk-test-key" assert headers["Content-Type"] == "application/json" + + +PDF_BASE64: Final = "JVBERi0xLjcK" +PNG_BASE64: Final = "iVBORw0KGgo=" + +ANTHROPIC_AGENT_TURN: Final[list[dict[str, object]]] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the indemnity terms."}, + { + "type": "document", + "title": "contract.pdf", + "source": {"type": "base64", "media_type": "application/pdf", "data": PDF_BASE64}, + }, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG_BASE64}}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me pull clause 14."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup_clause", "input": {"id": "14"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}], + }, +] + +ANTHROPIC_TOOLS: Final[list[dict[str, object]]] = [ + { + "name": "lookup_clause", + "description": "Fetch a clause by id", + "input_schema": {"type": "object", "properties": {"id": {"type": "string"}}}, + } +] + + +@pytest.mark.parametrize( + "system", + [None, "You review contracts.", [{"type": "text", "text": "You review contracts."}]], +) +def test_anthropic_messages_convert_exactly_as_v1_messages_inference_converts_them(system: object) -> None: + inference_request: Final[dict[str, object]] = { + "model": "gpt-4o", + "messages": ANTHROPIC_AGENT_TURN, + "tools": ANTHROPIC_TOOLS, + **({"system": system} if system is not None else {}), + } + inference: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_request( + inference_request, include_encrypted_reasoning=False + ) + + counted: Final = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=system, tools=ANTHROPIC_TOOLS + ) + + assert counted == ( + tuple(inference["input"]), + inference.get("instructions"), + tuple(inference["tools"]), + inference.get("text"), + ) + + +def test_anthropic_messages_keep_every_block_the_chat_converter_drops() -> None: + input_items, _, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=None, tools=None + ) + + assert input_items == ( + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize the indemnity terms."}, + { + "type": "input_file", + "filename": "contract.pdf", + "file_data": f"data:application/pdf;base64,{PDF_BASE64}", + }, + {"type": "input_image", "image_url": f"data:image/png;base64,{PNG_BASE64}"}, + ], + }, + {"type": "function_call", "call_id": "toolu_01", "name": "lookup_clause", "arguments": '{"id": "14"}'}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Let me pull clause 14."}], + }, + {"type": "function_call_output", "call_id": "toolu_01", "output": "14. Indemnification."}, + ) + + +def test_anthropic_list_system_becomes_joined_instruction_text() -> None: + _, instructions, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + system=[{"type": "text", "text": "First rule."}, {"type": "text", "text": "Second rule."}], + tools=None, + ) + + assert instructions == "First rule.\nSecond rule." diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index e4832a7c23b..3b4a1cfbf0c 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -30,7 +30,6 @@ from litellm.proxy.proxy_server import token_counter from litellm.types.utils import TokenCountResponse - @pytest.mark.asyncio async def test_vLLM_token_counting(): """ @@ -63,9 +62,7 @@ async def test_vLLM_token_counting(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the default tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the default tokenizer assert response.model_used == "wolfram/miquliz-120b-v2.0" @@ -98,9 +95,7 @@ async def test_token_counting_model_not_in_model_list(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the OpenAI tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the OpenAI tokenizer assert response.model_used == "special-alias" @@ -133,9 +128,7 @@ async def test_gpt_token_counting(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the OpenAI tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the OpenAI tokenizer assert response.request_model == "gpt-4" @@ -173,10 +166,8 @@ async def test_anthropic_messages_count_tokens_endpoint(): anthropic_endpoints._read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=False): - assert ( - call_endpoint == True - ), "Should be called with call_endpoint=True for Anthropic endpoint" + async def mock_token_counter(request, call_endpoint=False, messages_format="openai", output_format=None): + assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "claude-3-sonnet-20240229" assert request.messages == [{"role": "user", "content": "Hello Claude!"}] @@ -192,8 +183,8 @@ async def test_anthropic_messages_count_tokens_endpoint(): # Patch the imported token_counter function from proxy_server import litellm.proxy.proxy_server as proxy_server - original_token_counter = proxy_server.token_counter - proxy_server.token_counter = mock_token_counter + original_token_counter = proxy_server.count_request_tokens + proxy_server.count_request_tokens = mock_token_counter try: # Call the endpoint @@ -210,7 +201,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): finally: # Restore original functions anthropic_endpoints._read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -247,10 +238,8 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): anthropic_endpoints._read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=True): - assert ( - call_endpoint == True - ), "Should be called with call_endpoint=True for Anthropic endpoint" + async def mock_token_counter(request, call_endpoint=True, messages_format="openai", output_format=None): + assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "gpt-4" assert request.messages == [{"role": "user", "content": "Hello GPT!"}] @@ -266,8 +255,8 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): # Patch the imported token_counter function from proxy_server import litellm.proxy.proxy_server as proxy_server - original_token_counter = proxy_server.token_counter - proxy_server.token_counter = mock_token_counter + original_token_counter = proxy_server.count_request_tokens + proxy_server.count_request_tokens = mock_token_counter try: # Call the endpoint @@ -284,7 +273,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): finally: # Restore original functions anthropic_endpoints._read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -407,9 +396,7 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -466,9 +453,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -527,9 +512,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -544,9 +527,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): mock_router.model_list = [ { "model_name": "claude-3-5-sonnet", - "litellm_params": { - "model": "anthropic/claude-3-5-sonnet-20241022" - }, + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}, "model_info": {}, } ] @@ -555,9 +536,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): mock_router.async_get_available_deployment = AsyncMock( return_value={ "model_name": "claude-3-5-sonnet", - "litellm_params": { - "model": "anthropic/claude-3-5-sonnet-20241022" - }, + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}, "model_info": {}, } ) @@ -592,9 +571,7 @@ async def test_factory_registration(): assert counter is not None # Create test deployments - anthropic_deployment = { - "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"} - } + anthropic_deployment = {"litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}} non_anthropic_deployment = {"litellm_params": {"model": "openai/gpt-4"}} @@ -630,9 +607,7 @@ async def test_bedrock_count_tokens_endpoint(): model_list=[ { "model_name": "claude-bedrock", - "litellm_params": { - "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0" - }, + "litellm_params": {"model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0"}, } ] ) @@ -646,9 +621,7 @@ async def test_bedrock_count_tokens_endpoint(): } # Test the mock handler directly to verify correct parameter extraction - await mock_count_tokens_handler( - request_data, {}, "anthropic.claude-3-sonnet-20240229-v1:0" - ) + await mock_count_tokens_handler(request_data, {}, "anthropic.claude-3-sonnet-20240229-v1:0") @pytest.mark.asyncio @@ -710,10 +683,7 @@ async def test_vertex_ai_anthropic_token_counting(): assert call_args is not None assert call_args.kwargs["model"] == "claude-3-5-sonnet-20241022" assert "messages" in call_args.kwargs["request_data"] - assert ( - call_args.kwargs["request_data"]["messages"][0]["content"] - == "Hello Claude on Vertex AI! How are you?" - ) + assert call_args.kwargs["request_data"]["messages"][0]["content"] == "Hello Claude on Vertex AI! How are you?" # Validate response structure assert response.model_used == "claude-3-5-sonnet-20241022" @@ -746,9 +716,7 @@ def test_vertex_ai_partner_models_token_counting_endpoint(vertex_location): if vertex_location == "global": assert endpoint.startswith("https://aiplatform.googleapis.com") else: - assert endpoint.startswith( - f"https://{vertex_location}-aiplatform.googleapis.com" - ) + assert endpoint.startswith(f"https://{vertex_location}-aiplatform.googleapis.com") class _FailingRuntimeHandler(BedrockCountTokensHandler): @@ -834,20 +802,14 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): with patch.object(handler, "validate_count_tokens_request"): with patch.object(handler, "_get_aws_region_name", return_value="us-west-2"): - with patch.object( - handler, "transform_anthropic_to_bedrock_count_tokens", return_value={} - ): + with patch.object(handler, "transform_anthropic_to_bedrock_count_tokens", return_value={}): with patch.object( handler, "get_bedrock_count_tokens_endpoint", return_value="https://example.com", ): - with patch.object( - handler, "_sign_request", return_value=({}, "{}") - ): - with patch( - "litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client" - ) as mock_client: + with patch.object(handler, "_sign_request", return_value=({}, "{}")): + with patch("litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client") as mock_client: mock_async_client = AsyncMock() mock_async_client.post = AsyncMock(side_effect=http_error) mock_client.return_value = mock_async_client @@ -856,9 +818,7 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): await handler.handle_count_tokens_request( request_data={ "model": "test", - "messages": [ - {"role": "user", "content": "hello"} - ], + "messages": [{"role": "user", "content": "hello"}], }, litellm_params={}, resolved_model="anthropic.claude-3-sonnet", @@ -866,10 +826,7 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): assert exc_info.value.status_code == 403 # Message should be the raw response text - assert ( - exc_info.value.message - == "Forbidden - Invalid credentials" - ) + assert exc_info.value.message == "Forbidden - Invalid credentials" @pytest.mark.asyncio @@ -897,9 +854,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception(): mock_counter.count_tokens = AsyncMock(side_effect=http_error) # Save originals - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter original_router = litellm.proxy.proxy_server.llm_router try: @@ -907,9 +862,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception(): def mock_get_provider_token_counter(deployment, model_to_use): return (mock_counter, "claude-4-6-sonnet", "vertex_ai") - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter mock_router = MagicMock() mock_router.async_get_available_deployment = AsyncMock( @@ -937,9 +890,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception(): assert exc_info.value.type == "token_counting_error" assert exc_info.value.param == "model" finally: - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter litellm.proxy.proxy_server.llm_router = original_router @@ -975,9 +926,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): # Save original value and function original_disable = litellm.disable_token_counter - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter try: litellm.disable_token_counter = True @@ -991,9 +940,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): def mock_get_provider_token_counter(deployment, model_to_use): return (mock_counter, "anthropic.claude-3-sonnet", "bedrock") - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter with pytest.raises(ProxyException) as exc_info: await token_counter( @@ -1008,9 +955,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): assert "Rate limit exceeded" in exc_info.value.message finally: litellm.disable_token_counter = original_disable - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter @pytest.mark.asyncio @@ -1045,9 +990,7 @@ async def test_proxy_token_counter_error_raises_when_enabled(): # Save original value and function original_disable = litellm.disable_token_counter - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter try: litellm.disable_token_counter = False @@ -1061,9 +1004,7 @@ async def test_proxy_token_counter_error_raises_when_enabled(): def mock_get_provider_token_counter(deployment, model_to_use): return (mock_counter, "anthropic.claude-3-sonnet", "bedrock") - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter with patch("litellm.token_counter") as mock_local_counter: with pytest.raises(ProxyException) as exc_info: @@ -1081,9 +1022,7 @@ async def test_proxy_token_counter_error_raises_when_enabled(): mock_local_counter.assert_not_called() finally: litellm.disable_token_counter = original_disable - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter @pytest.mark.asyncio @@ -1109,10 +1048,10 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): original_read_request_body = anthropic_endpoints._read_request_body anthropic_endpoints._read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with Bedrock-style error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message='{"detail":{"message":"Input is too long for requested model."}}', type="token_counting_error", @@ -1120,7 +1059,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): code=400, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: with pytest.raises(HTTPException) as exc_info: @@ -1136,7 +1075,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): assert detail["error"]["message"] == "Input is too long for requested model." finally: anthropic_endpoints._read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -1161,10 +1100,10 @@ async def test_anthropic_endpoint_403_permission_error_format(): original_read_request_body = anthropic_endpoints._read_request_body anthropic_endpoints._read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 403 error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message='{"Message":"Bearer Token has expired"}', type="token_counting_error", @@ -1172,7 +1111,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): code=403, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: with pytest.raises(HTTPException) as exc_info: @@ -1186,7 +1125,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): assert detail["error"]["message"] == "Bearer Token has expired" finally: anthropic_endpoints._read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -1211,10 +1150,10 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): original_read_request_body = anthropic_endpoints._read_request_body anthropic_endpoints._read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 429 error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message="Rate limit exceeded", type="token_counting_error", @@ -1222,7 +1161,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): code=429, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: with pytest.raises(HTTPException) as exc_info: @@ -1236,4 +1175,4 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): assert detail["error"]["message"] == "Rate limit exceeded" finally: anthropic_endpoints._read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter diff --git a/tests/unit/test_count_tokens_public_api.py b/tests/unit/test_count_tokens_public_api.py index 84121dd63e4..e33bc1e5907 100644 --- a/tests/unit/test_count_tokens_public_api.py +++ b/tests/unit/test_count_tokens_public_api.py @@ -97,7 +97,7 @@ def test_acount_tokens_with_tools(): assert result.total_tokens == 30 mock_handler.assert_called_once() call_kwargs = mock_handler.call_args - assert call_kwargs.kwargs.get("tools") == tools + assert list(call_kwargs.kwargs.get("tools") or ()) == tools def test_acount_tokens_with_system(): From 350def8ccf349ad34b6229710f52fef8a56088d5 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:34:15 -0700 Subject: [PATCH 15/36] fix(count_tokens): count Anthropic-format requests for Gemini the way /v1/messages inference converts them (#46767) Backport of #46767 to stable/1.104.x. Cherry-picked from dba86bec28cb483c9937470c5d950f5e03699d3b (main). Co-authored-by: Mrinal Chanshetty --- .../adapters/transformation.py | 2 +- litellm/llms/gemini/common_utils.py | 26 +- litellm/llms/vertex_ai/common_utils.py | 25 +- .../llms/vertex_ai/count_tokens/context.py | 114 ++++++- .../llms/vertex_ai/gemini/transformation.py | 8 +- .../test_vertex_gemini_count_tokens_wire.py | 87 +++++ .../llms/gemini/test_gemini_common_utils.py | 98 ++++-- .../vertex_ai/count_tokens/test_context.py | 315 +++++++++++++++++- .../vertex_ai/test_vertex_ai_common_utils.py | 92 ++++- 9 files changed, 704 insertions(+), 63 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 85431a5a637..822f3081c5b 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -460,7 +460,7 @@ class LiteLLMAnthropicMessagesAdapter: *, custom_llm_provider: str | None = None, preserve_midturn_system: bool = False, - ) -> list: + ) -> list[AllMessageValues]: new_messages: Final[list[AllMessageValues]] = [] replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) leading_count: Final = next( diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index bd155c00c8f..e38160965f5 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -501,25 +501,41 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, count_context_error_response, + gemini_count_contents, + inference_request_view, + provider_error_response, vertex_gemini_count_context, ) - count_context: Final = vertex_gemini_count_context(system=system, tools=tools) + request_view: Final = inference_request_view(messages or [], system, messages_format) + count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) + resolved_contents: Final = ( + contents + if contents is not None + else await gemini_count_contents( + request_view.messages, messages_format, model_to_use, "gemini", count_tokens_params_request + ) + ) + if isinstance(resolved_contents, VertexGeminiCountContextError): + return count_context_error_response(resolved_contents, request_model, model_to_use) count_tokens_params: Final = { "model": model_to_use, - "contents": contents, + "contents": resolved_contents, "system_instruction": count_context.system_instruction, "tools": count_context.tools, } count_tokens_params_request.update(count_tokens_params) - result: Final = await GoogleAIStudioTokenCounter().acount_tokens( - **count_tokens_params_request, - ) + try: + result: Final = await GoogleAIStudioTokenCounter().acount_tokens( + **count_tokens_params_request, + ) + except litellm.APIError as error: + return provider_error_response(error, request_model, model_to_use) if result is not None: return TokenCountResponse( diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 7d4b87d0228..62a2e134f66 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1313,24 +1313,27 @@ class VertexAITokenCounter(BaseTokenCounter): from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, count_context_error_response, + gemini_count_contents, + inference_request_view, + provider_error_response, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, # pyright: ignore[reportPrivateUsage] # shared helper already used by gemini/chat, context_caching, and vertex_and_google_ai_studio_gemini - ) - count_context: Final = vertex_gemini_count_context(system=system, tools=tools) + request_view: Final = inference_request_view(messages or [], system, messages_format) + count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) resolved_contents: Final = ( contents if contents is not None - else _gemini_convert_messages_with_history( - messages=messages or [] # mutable-ok: fallback for None messages; helper signature requires list + else await gemini_count_contents( + request_view.messages, messages_format, model_to_use, "vertex_ai", count_tokens_params_request ) ) + if isinstance(resolved_contents, VertexGeminiCountContextError): + return count_context_error_response(resolved_contents, request_model, model_to_use) count_tokens_params: Final = { "model": model_to_use, @@ -1344,15 +1347,7 @@ class VertexAITokenCounter(BaseTokenCounter): **count_tokens_params_request, ) except litellm.APIError as error: - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="gemini", - error=True, - error_message=error.message, - status_code=error.status_code, - ) + return provider_error_response(error, request_model, model_to_use) if "totalTokens" in result: return TokenCountResponse( diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index 24ffb406a47..8377e5e00b5 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -1,12 +1,20 @@ -from collections.abc import Mapping +import copy +from collections.abc import Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Final +from typing import TYPE_CHECKING, Final, Literal from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never from litellm.types.llms.anthropic import AnthropicMessagesTextParam -from litellm.types.utils import TokenCountResponse +from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.vertex_ai import ContentType +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse + +if TYPE_CHECKING: + from litellm.exceptions import APIError + from litellm.llms.anthropic.experimental_pass_through.context_management.result import PolyfillResult @dataclass(frozen=True, slots=True) @@ -142,3 +150,103 @@ def count_context_error_response( error_message=error.error_message, status_code=400, ) + + +def provider_error_response(error: "APIError", request_model: str, model_used: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_used, + tokenizer_type="gemini", + error=True, + error_message=error.message, + status_code=error.status_code, + ) + + +@dataclass(frozen=True, slots=True) +class InferenceRequestView: + messages: Sequence[Mapping[str, object]] + system: object | None + + +def _client_compaction(messages: Sequence[Mapping[str, object]], system: object | None) -> "PolyfillResult | None": + from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import ( + apply_client_compaction_block_history, + ) + + chat_messages: Final = [dict(message) for message in messages] + if system is None or isinstance(system, str): + return apply_client_compaction_block_history(messages=chat_messages, system=system) + try: + blocks: Final = _SYSTEM_BLOCKS.validate_python(system) + except ValidationError: + return None + return apply_client_compaction_block_history(messages=chat_messages, system=[dict(block) for block in blocks]) + + +def inference_request_view( + messages: Sequence[Mapping[str, object]], system: object | None, messages_format: TokenCountMessagesFormat +) -> InferenceRequestView: + """Apply the history rewrites /v1/messages runs before it converts a body for Gemini.""" + from litellm.llms.anthropic.experimental_pass_through.messages.utils import normalize_native_messages + + match messages_format: + case "openai": + return InferenceRequestView(messages=messages, system=system) + case "anthropic": + normalized: Final = normalize_native_messages(messages) + compacted: Final = _client_compaction(normalized, system) + if compacted is None: + return InferenceRequestView(messages=normalized, system=system) + return InferenceRequestView(messages=compacted.messages, system=compacted.system) + case _: + assert_never(messages_format) + + +def _chat_messages( + messages: Sequence[Mapping[str, object]], + messages_format: TokenCountMessagesFormat, + model: str, + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> Sequence[AllMessageValues]: + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter + + match messages_format: + case "anthropic": + return LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=list(messages), # pyright: ignore[reportArgumentType] # raw Anthropic body, read the way /v1/messages inference reads it + model=model, + custom_llm_provider=custom_llm_provider, + ) + case "openai": + return copy.deepcopy(list(messages)) # pyright: ignore[reportReturnType] # chat messages from the request, copied because the AI Studio conversion rewrites them in place + case _: + assert_never(messages_format) + + +async def gemini_count_contents( + messages: Sequence[Mapping[str, object]], + messages_format: TokenCountMessagesFormat, + model: str, + custom_llm_provider: Literal["vertex_ai", "gemini"], + litellm_params: Mapping[str, object], +) -> Sequence[ContentType] | VertexGeminiCountContextError: + from litellm.exceptions import BadRequestError + from litellm.litellm_core_utils.asyncify import asyncify + from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media + from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig + from litellm.llms.vertex_ai.gemini.transformation import should_inline_remote_media + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + config: Final = GoogleAIStudioGeminiConfig() if custom_llm_provider == "gemini" else VertexGeminiConfig() + try: + chat_messages: Final = await async_inline_remote_media( + list(_chat_messages(messages, messages_format, model, custom_llm_provider)), + should_inline=should_inline_remote_media(custom_llm_provider), + ) + return await asyncify(config._transform_messages)( # pyright: ignore[reportPrivateUsage] # the message transform inference runs + messages=chat_messages, model=model, litellm_params=dict(litellm_params) + ) + except BadRequestError as error: + return VertexGeminiCountContextError(error_message=error.message) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index e3cc3bbb2dc..83b796d5df2 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -7,7 +7,7 @@ Why separate file? Make it easy to see how transformation works import json import os import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote @@ -1327,6 +1327,10 @@ def _vertex_inlines(media: RemoteMedia) -> bool: ) +def should_inline_remote_media(custom_llm_provider: str) -> Callable[[RemoteMedia], bool]: + return _ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + + async def async_transform_request_body( gemini_api_key: str | None, messages: list[AllMessageValues], @@ -1369,7 +1373,7 @@ async def async_transform_request_body( ) inlined_messages: Final = await async_inline_remote_media( - messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + messages, should_inline=should_inline_remote_media(custom_llm_provider) ) if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 1a148b19f8d..7fef79845a1 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -167,6 +167,93 @@ def test_messages_and_count_tokens_send_matching_gemini_context( assert count_body == _count_tokens_request_from_generate_content(generated_body) +_LOOKUP_TOOL: Final[JsonValue] = { + "name": "lookup", + "description": "Look up a value", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, +} +_CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "base64-pdf": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the document."}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + }, + ], + } + ], + }, + "base64-image": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + }, + "tool-turns": { + "tools": [_LOOKUP_TOOL], + "messages": [ + {"role": "user", "content": "Look up the value."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Looking it up."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "value"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}], + }, + ], + }, + } +) + + +def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key") + + +def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: + if "contents" in count_body: + return count_body["contents"] + return _JSON_OBJECT.validate_python(count_body["generateContentRequest"])["contents"] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +@pytest.mark.parametrize("case", sorted(_CONTENT_BLOCK_BODIES)) +def test_count_tokens_sends_the_contents_messages_inference_sends(gateway: Gateway, provider: str, case: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + body: Final = {"model": model, **_CONTENT_BLOCK_BODIES[case]} + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + requests: Final = peer.drain() + generated_body: Final = _json_body( + next(request for request in requests if request.target.endswith(":generateContent")) + ) + count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens"))) + assert _counted_contents(count_body) == generated_body["contents"] + + def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None: peer: Final = _GeminiPeer() with wire_server(peer) as wire, gateway.scenario() as scenario: diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 9a47d01ec46..e54023c6596 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -5,6 +5,7 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter +import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.common_utils import GeminiModelInfo, GoogleAIStudioTokenCounter from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler @@ -100,6 +101,58 @@ async def test_google_ai_studio_count_tokens_rejects_image_system_without_reques assert peer.request_body is None +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_rejects_an_image_url_it_refuses_to_fetch_without_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}], + } + ], + contents=None, + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "http://127.0.0.1:9/image.png" in (result.error_message or "") + assert peer.request_body is None + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_answers_with_the_status_google_rejected_it_with() -> None: + rejection: Final = '{"error": {"code": 400, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}' + client: Final = AsyncHTTPHandler( + transport=httpx.MockTransport(lambda request: httpx.Response(400, text=rejection, request=request)) + ) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "Invalid function name" in (result.error_message or "") + + @pytest.mark.asyncio @pytest.mark.parametrize( ("system_instruction", "tools", "expected_body"), @@ -129,9 +182,7 @@ async def test_gemini_count_tokens_sends_optional_request_fields( expected_body: dict[str, JsonValue], ) -> None: peer: Final = _GeminiCountTokensPeer() - token_counter: Final = GeminiCountTokensHandler( - client=AsyncHTTPHandler(transport=httpx.MockTransport(peer)) - ) + token_counter: Final = GeminiCountTokensHandler(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer))) result: Final = await token_counter.acount_tokens( contents=[{"role": "user", "parts": [{"text": "hello"}]}], @@ -177,18 +228,10 @@ class TestGeminiModelInfo: # Test edge cases where model names end with characters from "models/" # These would be incorrectly processed if using strip("models/") instead of replace("models/", "") models = [ - { - "name": "models/gemini-1.5-pro" - }, # ends with 'o' - would become "gemini-1.5-pr" with strip() - { - "name": "models/test-model" - }, # ends with 'l' - would become "gemini/test-mode" with strip() - { - "name": "models/custom-models" - }, # ends with 's' - would become "gemini/custom-model" with strip() - { - "name": "models/demo" - }, # ends with 'o' - would become "gemini/dem" with strip() + {"name": "models/gemini-1.5-pro"}, # ends with 'o' - would become "gemini-1.5-pr" with strip() + {"name": "models/test-model"}, # ends with 'l' - would become "gemini/test-mode" with strip() + {"name": "models/custom-models"}, # ends with 's' - would become "gemini/custom-model" with strip() + {"name": "models/demo"}, # ends with 'o' - would become "gemini/dem" with strip() ] result = gemini_model_info.process_model_name(models) @@ -239,16 +282,10 @@ class TestGoogleAIStudioTokenCounter: token_counter = GoogleAIStudioTokenCounter() # Test with gemini provider - should return True - assert ( - token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) - is True - ) + assert token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) is True # Test with other providers - should return False - assert ( - token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) - is False - ) + assert token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) is False assert token_counter.should_use_token_counting_api("anthropic") is False assert token_counter.should_use_token_counting_api("vertex_ai") is False @@ -320,9 +357,7 @@ class TestGoogleAIStudioTokenCounter: "functionResponse": { "id": "read_many_files-1757526647518-730a691aac11c", # This should be removed "name": "read_many_files", - "response": { - "output": "No files matching the criteria were found or all were skipped." - }, + "response": {"output": "No files matching the criteria were found or all were skipped."}, } } ], @@ -331,9 +366,7 @@ class TestGoogleAIStudioTokenCounter: ] # Clean the contents - cleaned_contents = token_counter._clean_contents_for_gemini_api( - contents_with_id - ) + cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_with_id) # Verify the 'id' field was removed function_response = cleaned_contents[1]["parts"][0]["functionResponse"] @@ -342,8 +375,7 @@ class TestGoogleAIStudioTokenCounter: assert "response" in function_response assert function_response["name"] == "read_many_files" assert ( - function_response["response"]["output"] - == "No files matching the criteria were found or all were skipped." + function_response["response"]["output"] == "No files matching the criteria were found or all were skipped." ) def test_clean_contents_for_gemini_api_preserves_other_fields(self): @@ -359,9 +391,7 @@ class TestGoogleAIStudioTokenCounter: ] # Clean the contents - cleaned_contents = token_counter._clean_contents_for_gemini_api( - contents_without_function_response - ) + cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_without_function_response) # Verify the contents are unchanged assert cleaned_contents == contents_without_function_response diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py index 1c05a814fa4..4929b2c8216 100644 --- a/tests/unit/llms/vertex_ai/count_tokens/test_context.py +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -1,9 +1,19 @@ +import asyncio +import base64 +import time +import uuid +from collections.abc import Generator +from contextlib import contextmanager from copy import deepcopy -from typing import Final +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Thread +from typing import Final, Literal +from unittest.mock import patch import pytest -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter +import litellm from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, truncate_tool_name, @@ -11,6 +21,8 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContext, VertexGeminiCountContextError, + gemini_count_contents, + inference_request_view, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig @@ -194,3 +206,302 @@ def test_invalid_tools_return_an_error_value(tools: JsonValue) -> None: assert isinstance(context, VertexGeminiCountContextError) assert "tools" in context.error_message + + +_ANTHROPIC_CONVERSATION: Final = ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize both"}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0="}}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + ], + }, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "x"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}]}, +) + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_anthropic_documents_images_and_tool_turns_become_gemini_parts( + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> None: + contents: Final = await gemini_count_contents( + _ANTHROPIC_CONVERSATION, "anthropic", "gemini-2.5-flash", custom_llm_provider, {} + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Summarize both"}, + {"inline_data": {"data": "JVBERi0=", "mime_type": "application/pdf"}}, + {"inline_data": {"data": "aW1hZ2U=", "mime_type": "image/png"}}, + ], + }, + {"role": "model", "parts": [{"function_call": {"name": "lookup", "args": {"query": "x"}}}]}, + {"role": "user", "parts": [{"function_response": {"name": "lookup", "response": {"content": "42"}}}]}, + ] + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_openai_messages_keep_their_chat_conversion(custom_llm_provider: Literal["vertex_ai", "gemini"]) -> None: + contents: Final = await gemini_count_contents( + ({"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi"}), + "openai", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + + assert contents == [ + {"role": "user", "parts": [{"text": "hello"}]}, + {"role": "model", "parts": [{"text": "hi"}]}, + ] + + +_PNG_BYTES: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +) +_DOWNLOAD_SECONDS: Final = 0.5 + + +class _SlowPngHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + time.sleep(_DOWNLOAD_SECONDS) + self.send_response(200) + self.send_header("Content-Type", "image/png") + self.send_header("Content-Length", str(len(_PNG_BYTES))) + self.end_headers() + self.wfile.write(_PNG_BYTES) + + def log_message(self, format: str, *args: object) -> None: + return None + + +@contextmanager +def _slow_png_server() -> Generator[str, None, None]: + with ThreadingHTTPServer(("127.0.0.1", 0), _SlowPngHandler) as server: + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + + +async def _event_loop_wake_delay(nap_seconds: float) -> float: + started: Final = time.monotonic() + await asyncio.sleep(nap_seconds) + return time.monotonic() - started + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_remote_image_downloads_without_blocking_the_event_loop( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with _slow_png_server() as base_url: + image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png" + wake_delay, contents = await asyncio.gather( + _event_loop_wake_delay(0.05), + gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe it"}, + {"type": "image", "source": {"type": "url", "url": image_url}}, + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ), + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Describe it"}, + {"inline_data": {"data": base64.b64encode(_PNG_BYTES).decode(), "mime_type": "image/png"}}, + ], + } + ] + assert wake_delay < _DOWNLOAD_SECONDS / 2 + + +class _MissingImageHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: + return None + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_image_that_cannot_be_fetched_returns_an_error_value( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with ThreadingHTTPServer(("127.0.0.1", 0), _MissingImageHandler) as server: + Thread(target=server.serve_forever, daemon=True).start() + try: + contents = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "url", + "url": f"http://127.0.0.1:{server.server_address[1]}/{uuid.uuid4().hex}.png", + }, + } + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + finally: + server.shutdown() + + assert isinstance(contents, VertexGeminiCountContextError) + + +@pytest.mark.parametrize( + "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") +) +def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: + view: Final = inference_request_view( + ( + {"role": "user", "content": "first question " * 50}, + {"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]}, + {"role": "user", "content": "latest question"}, + ), + system, + "anthropic", + ) + + assert view.messages == [{"role": "user", "content": "latest question"}] + assert "Summary of the earlier turns." in str(view.system) + assert "Be brief." in str(view.system) + + +def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + view: Final = inference_request_view( + ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + ), + None, + "anthropic", + ) + + assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"]) + assert [block["type"] for block in assistant_blocks] == ["text"] + assert snippet in assistant_blocks[0]["text"] + + +def test_openai_view_is_the_request_unchanged() -> None: + messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) + + view: Final = inference_request_view(messages, "Be brief.", "openai") + + assert view.messages == messages + assert view.system == "Be brief." + + +async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None: + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value="image/png", + ) as lookup: + contents: Final = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "url", "url": "gs://bucket/image-without-extension"}} + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + "vertex_ai", + {"vertex_project": "deployment-project"}, + ) + + assert lookup.call_args.kwargs.get("vertex_project") == "deployment-project" + assert not isinstance(contents, VertexGeminiCountContextError) + assert contents[0]["parts"][0] == { + "file_data": {"mime_type": "image/png", "file_uri": "gs://bucket/image-without-extension"} + } + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +def test_anthropic_replayed_web_search_result_is_counted_as_its_text( + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + messages: Final = ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + {"role": "user", "content": "Thanks."}, + ) + + view: Final = inference_request_view(messages, None, "anthropic") + contents: Final = asyncio.run( + gemini_count_contents(view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {}) + ) + + assert snippet in str(contents) diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 59cd53567ae..f9561a0d710 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1,3 +1,5 @@ +from collections.abc import Iterator +from contextlib import contextmanager from typing import Final from unittest.mock import AsyncMock, patch @@ -5,6 +7,7 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter +import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler @@ -21,16 +24,21 @@ from litellm.llms.vertex_ai.common_utils import ( vertex_request_labels_from_litellm_params, ) from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler +from litellm.types.utils import TokenCountMessagesFormat @pytest.mark.asyncio -async def test_vertex_gemini_rejects_non_text_system_blocks_before_request() -> None: +@pytest.mark.parametrize("messages_format", ("openai", "anthropic")) +async def test_vertex_gemini_rejects_non_text_system_blocks_before_request( + messages_format: TokenCountMessagesFormat, +) -> None: result: Final = await VertexAITokenCounter().count_tokens( model_to_use="gemini-4-argon", messages=[{"role": "user", "content": "Count this"}], contents=None, request_model="vertex_ai/gemini-4-argon", system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + messages_format=messages_format, ) assert result is not None @@ -39,6 +47,88 @@ async def test_vertex_gemini_rejects_non_text_system_blocks_before_request() -> assert "text blocks" in (result.error_message or "") +class _VertexCountTokensPeer: + def __init__(self, status_code: int) -> None: + self.status_code: Final = status_code + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + if self.status_code == 200: + return httpx.Response(200, json={"totalTokens": 37}, request=request) + return httpx.Response( + self.status_code, + json={"error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}, + request=request, + ) + + +@contextmanager +def _vertex_count_tokens_peer(status_code: int) -> Iterator[_VertexCountTokensPeer]: + peer: Final = _VertexCountTokensPeer(status_code) + with ( + patch.object( + gemini_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ), + patch.object( + vertex_count_tokens_handler.VertexAITokenCounter, + "_ensure_access_token_async", + new=AsyncMock(return_value=("fake-token", "test-project")), + ), + ): + yield peer + + +_VERTEX_DEPLOYMENT: Final = {"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-central1"}} + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_answers_with_the_status_google_rejected_it_with() -> None: + with _vertex_count_tokens_peer(400) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "Invalid function name" in (result.error_message or "") + assert len(peer.requests) == 1 + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_rejects_an_image_url_it_refuses_to_fetch_without_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + with _vertex_count_tokens_peer(200) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}], + } + ], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "http://127.0.0.1:9/image.png" in (result.error_message or "") + assert peer.requests == () + + @pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"]) def test_validate_vertex_location_accepts_valid(location): assert validate_vertex_location(location) == location From 13b4537c81980184e8ffb07d8c5841122270aaa1 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:35:25 -0700 Subject: [PATCH 16/36] fix(count_tokens): return a 400 naming Anthropic blocks the OpenAI and Gemini count can't convert (#46772) Backport of #46772 to stable/1.104.x. Cherry-picked from 79ca820933a3db556759c14e3ee0a01763a4deaf (main). Co-authored-by: Mrinal Chanshetty --- .../convertible_blocks.py | 106 ++++++++++ litellm/llms/gemini/common_utils.py | 4 + .../responses/count_tokens/token_counter.py | 6 +- litellm/llms/vertex_ai/common_utils.py | 4 + .../llms/vertex_ai/count_tokens/context.py | 9 + .../proxy/anthropic_endpoints/endpoints.py | 22 +- .../test_openai_count_tokens_tools_wire.py | 42 ++++ .../test_vertex_gemini_count_tokens_wire.py | 47 +++++ .../test_convertible_blocks.py | 192 ++++++++++++++++++ tests/unit/proxy/test_proxy_token_counter.py | 37 ++-- 10 files changed, 430 insertions(+), 39 deletions(-) create mode 100644 litellm/llms/anthropic/experimental_pass_through/convertible_blocks.py create mode 100644 tests/unit/llms/anthropic/experimental_pass_through/test_convertible_blocks.py diff --git a/litellm/llms/anthropic/experimental_pass_through/convertible_blocks.py b/litellm/llms/anthropic/experimental_pass_through/convertible_blocks.py new file mode 100644 index 00000000000..0b86ccda012 --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/convertible_blocks.py @@ -0,0 +1,106 @@ +from collections.abc import Iterator, Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.llms.anthropic.experimental_pass_through.messages.utils import normalize_native_messages +from litellm.types.utils import TokenCountMessagesFormat + +_CONVERTIBLE_BLOCK_TYPES: Final = MappingProxyType( + { + "user": frozenset({"text", "image", "document", "tool_result"}), + "assistant": frozenset({"text", "tool_use", "thinking", "redacted_thinking"}), + } +) +_SOURCED_BLOCK_TYPES: Final = frozenset({"image", "document"}) +_TOOL_RESULT_BLOCK_TYPES: Final = frozenset({"text", "image", "document"}) + + +_JSON_MESSAGES: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) + + +def _json_messages(messages: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, JsonValue], ...]: + try: + return _JSON_MESSAGES.validate_python(messages) + except ValidationError: + return () + + +def _source_label(block_type: str, source: JsonValue) -> str | None: + if not isinstance(source, dict): + return f"{block_type} without a source" + match source.get("type"): + case "base64" if source.get("data"): + return None + case "url" if source.get("url"): + return None + case "base64" | "url" as source_type: + return f"{block_type} with an empty {source_type} source" + case str() as source_type: + return f"{block_type} with a {source_type} source" + case _: + return f"{block_type} with an untyped source" + + +def _block_label(block: JsonValue, convertible: frozenset[str]) -> str | None: + if not isinstance(block, dict): + return None + block_type: Final = block.get("type") + if not isinstance(block_type, str): + return "a block without a type" + if block_type not in convertible: + return block_type + return _source_label(block_type, block.get("source")) if block_type in _SOURCED_BLOCK_TYPES else None + + +def _tool_result_labels(block: JsonValue) -> Iterator[str]: + if not isinstance(block, dict) or block.get("type") != "tool_result": + return iter(()) + inner: Final = block.get("content") + if not isinstance(inner, list): + return iter(()) + return ( + f"{label} inside a tool_result" + for item in inner + if (label := _block_label(item, _TOOL_RESULT_BLOCK_TYPES)) is not None + ) + + +def _content_labels(block: JsonValue, convertible: frozenset[str]) -> Iterator[str]: + label: Final = _block_label(block, convertible) + return chain(() if label is None else (label,), _tool_result_labels(block)) + + +def _message_labels(message: Mapping[str, JsonValue]) -> Iterator[str]: + role: Final = message.get("role") + convertible: Final = _CONVERTIBLE_BLOCK_TYPES.get(role) if isinstance(role, str) else None + content: Final = message.get("content") + if convertible is None or not isinstance(content, list): + return iter(()) + return chain.from_iterable(_content_labels(block, convertible) for block in content) + + +def unconvertible_anthropic_blocks(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + normalized: Final = _json_messages(normalize_native_messages(messages)) + return tuple(sorted(set(chain.from_iterable(_message_labels(message) for message in normalized)))) + + +def unconvertible_blocks_error( + messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat +) -> str | None: + match messages_format: + case "openai": + return None + case "anthropic": + labels: Final = unconvertible_anthropic_blocks(messages or ()) + return ( + "count_tokens can't convert these content blocks for this model, so the count would leave them out: " + + ", ".join(labels) + if labels + else None + ) + case _: + assert_never(messages_format) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index e38160965f5..155366f476e 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -504,6 +504,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): gemini_count_contents, inference_request_view, provider_error_response, + unconvertible_messages_error, vertex_gemini_count_context, ) @@ -511,6 +512,9 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) + blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format) + if blocks_error is not None: + return count_context_error_response(blocks_error, request_model, model_to_use) deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 59c812a0a32..14f3884cc71 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -10,6 +10,7 @@ from pydantic import JsonValue from typing_extensions import assert_never from litellm._logging import verbose_logger +from litellm.llms.anthropic.experimental_pass_through.convertible_blocks import unconvertible_blocks_error from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) @@ -91,14 +92,15 @@ class OpenAITokenCounter(BaseTokenCounter): if messages_format == "anthropic" else None ) - if malformed_error is not None: + request_error: Final = malformed_error or unconvertible_blocks_error(messages, messages_format) + if request_error is not None: return TokenCountResponse( total_tokens=0, request_model=request_model, model_used=model_to_use, tokenizer_type="openai_api", error=True, - error_message=malformed_error, + error_message=request_error, status_code=400, ) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 62a2e134f66..2c7e07616b5 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1316,6 +1316,7 @@ class VertexAITokenCounter(BaseTokenCounter): gemini_count_contents, inference_request_view, provider_error_response, + unconvertible_messages_error, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter @@ -1324,6 +1325,9 @@ class VertexAITokenCounter(BaseTokenCounter): count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) + blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format) + if blocks_error is not None: + return count_context_error_response(blocks_error, request_model, model_to_use) resolved_contents: Final = ( contents diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index 8377e5e00b5..3c9c30b4a29 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -138,6 +138,15 @@ def vertex_gemini_count_context( ) +def unconvertible_messages_error( + messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat +) -> VertexGeminiCountContextError | None: + from litellm.llms.anthropic.experimental_pass_through.convertible_blocks import unconvertible_blocks_error + + error_message: Final = unconvertible_blocks_error(messages, messages_format) + return None if error_message is None else VertexGeminiCountContextError(error_message=error_message) + + def count_context_error_response( error: VertexGeminiCountContextError, request_model: str, model_used: str ) -> TokenCountResponse: diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index f166fe87b4a..2622446a699 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -40,7 +40,6 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) -from litellm.types.utils import TokenCountResponse router: Final = APIRouter() @@ -358,30 +357,15 @@ async def count_tokens( messages_format="anthropic", output_format=_requested_output_format(data), ) - _token_response_dict: dict = {} - if isinstance(token_response, TokenCountResponse): - _token_response_dict = token_response.model_dump() - elif isinstance(token_response, dict): - _token_response_dict = token_response - - # Convert the internal response to Anthropic API format - return {"input_tokens": _token_response_dict.get("total_tokens", 0)} - except HTTPException: raise except ProxyException as e: - status_code: Final = int(e.code) if e.code and e.code.isdigit() else 500 - detail: Final = AnthropicExceptionMapping.transform_to_anthropic_error( - status_code=status_code, - raw_message=e.message, - ) - raise HTTPException( - status_code=status_code, - detail=detail, - ) + return _anthropic_error_json_response(e, request) except Exception as e: log_llm_api_exception(e, litellm_call_id) raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) + else: + return {"input_tokens": token_response.total_tokens} @router.post( diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index 3ab5826006f..2b23ccdf4f2 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -321,6 +321,48 @@ def test_messages_count_tokens_sends_the_input_v1_messages_inference_sends( assert count_body["model"] == inference_body["model"] +_UNCONVERTIBLE_BODY: Final[Mapping[str, JsonValue]] = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}}, + {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "search_result", + "source": "https://example.com/hits", + "title": "Hits", + "content": [{"type": "text", "text": "hit"}], + } + ], + } + ], + }, + ], +} + + +def test_messages_count_tokens_rejects_blocks_it_cannot_convert(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, **_UNCONVERTIBLE_BODY}) + assert response.status_code == 400, response.text + assert "document with a text source, search_result, search_result inside a tool_result" in response.text + assert wire.drain() == () + + def test_messages_count_tokens_without_a_role_answers_400_before_calling_openai(counting_proxy: OwnedProxy) -> None: gateway: Final = counting_proxy.gateway with wire_server(_count_response) as wire, gateway.scenario() as scenario: diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 7fef79845a1..5fe82d37e9e 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -254,6 +254,53 @@ def test_count_tokens_sends_the_contents_messages_inference_sends(gateway: Gatew assert _counted_contents(count_body) == generated_body["contents"] +_UNCONVERTIBLE_MESSAGES: Final[JsonValue] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}}, + {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "search_result", + "source": "https://example.com/hits", + "title": "Hits", + "content": [{"type": "text", "text": "hit"}], + } + ], + } + ], + }, +] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_count_tokens_rejects_blocks_it_cannot_convert(gateway: Gateway, provider: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _UNCONVERTIBLE_MESSAGES} + ) + assert response.status_code == 400, response.text + assert "document with a text source, search_result, search_result inside a tool_result" in response.text + assert peer.drain() == () + + def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None: peer: Final = _GeminiPeer() with wire_server(peer) as wire, gateway.scenario() as scenario: diff --git a/tests/unit/llms/anthropic/experimental_pass_through/test_convertible_blocks.py b/tests/unit/llms/anthropic/experimental_pass_through/test_convertible_blocks.py new file mode 100644 index 00000000000..fe80a85d57b --- /dev/null +++ b/tests/unit/llms/anthropic/experimental_pass_through/test_convertible_blocks.py @@ -0,0 +1,192 @@ +import asyncio +import json +from collections.abc import Mapping +from typing import Final + +import pytest +from pydantic import JsonValue, TypeAdapter + +from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter +from litellm.llms.anthropic.experimental_pass_through.convertible_blocks import ( + unconvertible_anthropic_blocks, + unconvertible_blocks_error, +) +from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.llms.vertex_ai.count_tokens.context import gemini_count_contents +from litellm.types.llms.anthropic import AllAnthropicPassThroughMessageValues + +_JSON_OBJECTS: Final = TypeAdapter(list[dict[str, JsonValue]]) + +_PDF: Final = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"} +_PNG: Final = {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="} + +_CONVERTED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object]], ...]] = ( + ("user", {"type": "text", "text": "hi"}), + ("user", {"type": "image", "source": _PNG}), + ("user", {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}), + ("user", {"type": "document", "source": _PDF}), + ("user", {"type": "document", "source": {"type": "url", "url": "https://example.com/a.pdf"}}), + ("user", {"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}), + ("assistant", {"type": "text", "text": "hi"}), + ("assistant", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}), + ("assistant", {"type": "thinking", "thinking": "Let me think.", "signature": "sig"}), +) + +_DROPPED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object], str], ...]] = ( + ("user", {"type": "search_result", "source": "https://example.com", "title": "t", "content": []}, "search_result"), + ("user", {"type": "document", "source": {"type": "text", "data": "notes"}}, "document with a text source"), + ("user", {"type": "image", "source": {"type": "file", "file_id": "file_01"}}, "image with a file source"), + ("user", {"type": "image", "source": {"type": "base64", "data": ""}}, "image with an empty base64 source"), + ("user", {"type": "document"}, "document without a source"), + ("user", {"text": "untyped"}, "a block without a type"), + ("user", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}, "tool_use"), + ( + "assistant", + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {}}, + "server_tool_use", + ), +) + + +def _adapter_messages(role: str, block: Mapping[str, object]) -> list[AllAnthropicPassThroughMessageValues]: + return [{"role": role, "content": [dict(block)]}] # pyright: ignore[reportReturnType] # sends block shapes the TypedDicts leave out on purpose + + +def _converted_by_both_adapters(role: str, block: Mapping[str, object]) -> tuple[bool, bool]: + messages: Final = _adapter_messages(role, block) + chat: Final = _JSON_OBJECTS.validate_python( + LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=messages) + ) + responses: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input(messages) + chat_kept: Final = any( + message.get("content") or message.get("tool_calls") or message.get("thinking_blocks") for message in chat + ) + return chat_kept, bool(responses) + + +@pytest.mark.parametrize(("role", "block"), _CONVERTED_BLOCKS) +def test_blocks_both_conversions_keep_are_accepted(role: str, block: Mapping[str, object]) -> None: + assert _converted_by_both_adapters(role, block) == (True, True) + assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == () + + +@pytest.mark.parametrize(("role", "block", "label"), _DROPPED_BLOCKS) +def test_blocks_both_conversions_drop_are_named(role: str, block: Mapping[str, object], label: str) -> None: + assert _converted_by_both_adapters(role, block) == (False, False) + assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == (label,) + + +def test_labels_are_listed_once_in_order_and_string_content_is_accepted() -> None: + messages: Final = [ + {"role": "user", "content": "plain text"}, + {"role": "user", "content": [{"type": "search_result"}, {"type": "container_upload"}]}, + {"role": "assistant", "content": [{"type": "search_result"}]}, + ] + assert unconvertible_anthropic_blocks(messages) == ("container_upload", "search_result") + + +def test_redacted_thinking_is_accepted_like_other_replayed_thinking() -> None: + messages: Final = [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "opaque"}]}] + assert unconvertible_anthropic_blocks(messages) == () + + +def test_messages_that_are_not_json_are_left_to_the_count() -> None: + messages: Final = [{"role": "user", "content": [{"type": "search_result", "source": object()}]}] + assert unconvertible_anthropic_blocks(messages) == () + + +def _tool_turn(inner: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + return ( + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": [dict(inner)]}]}, + ) + + +def _sent_by_each_counter(inner: Mapping[str, object]) -> tuple[str, str]: + messages: Final = _tool_turn(inner) + gemini: Final = asyncio.run(gemini_count_contents(messages, "anthropic", "gemini-2.5-flash", "gemini", {})) + openai: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input( + [dict(message) for message in messages] # pyright: ignore[reportArgumentType] # sends block shapes the TypedDicts leave out on purpose + ) + return json.dumps(gemini, default=str), json.dumps(openai, default=str) + + +_KEPT_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str], ...]] = ( + ({"type": "text", "text": "forty-two"}, "forty-two"), + ({"type": "image", "source": _PNG}, _PNG["data"]), + ({"type": "document", "source": _PDF}, _PDF["data"]), +) + +_DROPPED_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str, str], ...]] = ( + ( + { + "type": "search_result", + "source": "https://example.com", + "title": "t", + "content": [{"type": "text", "text": "search-snippet"}], + }, + "search-snippet", + "search_result inside a tool_result", + ), + ( + {"type": "document", "source": {"type": "text", "data": "text-notes"}}, + "text-notes", + "document with a text source inside a tool_result", + ), + ({"type": "tool_reference", "tool_name": "find_docs"}, "find_docs", "tool_reference inside a tool_result"), +) + + +@pytest.mark.parametrize(("inner", "marker"), _KEPT_INSIDE_TOOL_RESULT) +def test_tool_result_content_both_counters_send_is_accepted(inner: Mapping[str, object], marker: str) -> None: + gemini, openai = _sent_by_each_counter(inner) + + assert marker in gemini + assert marker in openai + assert unconvertible_anthropic_blocks(_tool_turn(inner)) == () + + +@pytest.mark.parametrize(("inner", "marker", "label"), _DROPPED_INSIDE_TOOL_RESULT) +def test_tool_result_content_both_counters_drop_is_named(inner: Mapping[str, object], marker: str, label: str) -> None: + gemini, openai = _sent_by_each_counter(inner) + + assert marker not in gemini + assert marker not in openai + assert unconvertible_anthropic_blocks(_tool_turn(inner)) == (label,) + + +def _replayed_web_search(encrypted_content: str) -> tuple[Mapping[str, object], ...]: + return ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": encrypted_content, + "snippet": "LiteLLM ships a weekly release candidate.", + } + ], + }, + ], + }, + ) + + +def test_a_replayed_web_search_inference_sends_as_text_is_accepted() -> None: + assert unconvertible_blocks_error(_replayed_web_search(""), "anthropic") is None + + +def test_an_encrypted_web_search_inference_cannot_send_is_named() -> None: + error: Final = unconvertible_blocks_error(_replayed_web_search("opaque"), "anthropic") + assert error is not None + assert error.endswith("server_tool_use, web_search_tool_result") diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 3b4a1cfbf0c..c1d437cd886 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2,6 +2,8 @@ # 1. Generate a Key, and use it to make a call +import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -14,7 +16,8 @@ load_dotenv() # this file is to test litellm/proxy -from fastapi import HTTPException, Request +from fastapi import Request +from fastapi.responses import JSONResponse import litellm from litellm import Router @@ -1035,6 +1038,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1062,14 +1066,11 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - # Verify HTTP status code is correct - assert exc_info.value.status_code == 400 - - # Verify error is in Anthropic format - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "invalid_request_error" assert detail["error"]["message"] == "Input is too long for requested model." @@ -1087,6 +1088,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): import litellm.proxy.proxy_server as proxy_server mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1114,12 +1116,11 @@ async def test_anthropic_endpoint_403_permission_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 403 - - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 403 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "permission_error" assert detail["error"]["message"] == "Bearer Token has expired" @@ -1137,6 +1138,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): import litellm.proxy.proxy_server as proxy_server mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1164,12 +1166,11 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 429 - - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 429 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "rate_limit_error" assert detail["error"]["message"] == "Rate limit exceeded" From 4d5a6577c2c17077c6a8cc1ab85dfdda85e9dc8d Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:46:19 -0700 Subject: [PATCH 17/36] fix(proxy)!: count endpoints return errors in their source API shape instead of silent local counts (#46815) Backport of #46815 to stable/1.104.x. Cherry-picked from b21e659c24720946ed609c0669142aec9d0f2422 (main). Co-authored-by: Mrinal Chanshetty --- .../gitlab/gitlab_prompt_manager.py | 2 +- .../anthropic/count_tokens/token_counter.py | 23 +- .../anthropic/count_tokens/handler.py | 11 +- .../anthropic/count_tokens/token_counter.py | 33 ++- .../anthropic/count_tokens/transformation.py | 9 +- .../responses/count_tokens/token_counter.py | 61 ++++- .../proxy/anthropic_endpoints/endpoints.py | 11 +- litellm/proxy/common_request_processing.py | 21 ++ litellm/proxy/google_endpoints/endpoints.py | 71 +++++- .../test_count_tokens_error_shapes_wire.py | 205 +++++++++++++++++ .../anthropic_endpoints/test_endpoints.py | 7 +- .../test_google_api_endpoints.py | 23 ++ .../llms/anthropic/test_count_tokens_oauth.py | 51 +++++ .../count_tokens/test_token_counter.py | 212 ++++++++++++++++++ .../count_tokens/test_token_counter.py | 41 ++++ .../openai/test_openai_workload_identity.py | 53 +++++ tests/unit/proxy/test_proxy_token_counter.py | 45 ++-- 17 files changed, 800 insertions(+), 79 deletions(-) create mode 100644 tests/integration/providers/test_count_tokens_error_shapes_wire.py create mode 100644 tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index 817d280074f..ee1d4dc7043 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -215,7 +215,7 @@ class GitLabTemplateManager: elif value.replace(".", "").isdigit(): try: result[key] = float(value) - except Exception: + except ValueError: result[key] = value else: result[key] = value.strip("\"'") diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 1ca882071cb..2154189a9c7 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -7,6 +7,7 @@ from typing import Any, Final from pydantic import JsonValue +import litellm from litellm._logging import verbose_logger from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.base_llm.base_utils import BaseTokenCounter @@ -58,14 +59,22 @@ class AnthropicTokenCounter(BaseTokenCounter): deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) - # Get Anthropic API key from deployment config or environment - api_key = litellm_params.get("api_key") + api_key: Final = ( + litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key or os.getenv("ANTHROPIC_API_KEY") + ) if not api_key: - api_key = os.getenv("ANTHROPIC_API_KEY") - - if not api_key: - verbose_logger.warning("No Anthropic API key found for token counting") - return None + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="anthropic_api", + error=True, + error_message=( + f"No Anthropic API key is set for model {request_model or model_to_use}, so tokens can't be " + "counted. Set api_key on the deployment or ANTHROPIC_API_KEY in the proxy environment." + ), + status_code=401, + ) try: result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 5cbde9ef52b..e1eb83e4ebf 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -29,7 +29,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): self, model: str, messages: list[dict[str, Any]], - api_key: str, + api_key: str | None, api_base: str, litellm_params: dict[str, Any] | None = None, timeout: float | httpx.Timeout | None = None, @@ -79,6 +79,15 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): api_key=api_key, litellm_params=litellm_params, ) + if not {"x-api-key", "Authorization"} & headers.keys(): + raise AnthropicError( + status_code=401, + message=( + f"No Azure AI API key or Entra ID credential is set for model {model}, so tokens can't be " + "counted. Set api_key or an Entra ID credential on the deployment, or AZURE_AI_API_KEY or " + "AZURE_API_KEY in the proxy environment." + ), + ) # Use LiteLLM's async httpx client async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.AZURE_AI) 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 c02785f2051..5de5240ee47 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -7,6 +7,7 @@ from typing import Any, Final from pydantic import JsonValue +import litellm from litellm._logging import verbose_logger from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, @@ -18,6 +19,21 @@ from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCou azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler() +def _missing_api_base_response(request_model: str, model_to_use: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="azure_ai_anthropic_api", + error=True, + error_message=( + f"No Azure AI API base is set for model {request_model or model_to_use}, so tokens can't be counted. " + "Set api_base on the deployment or AZURE_AI_API_BASE or AZURE_API_BASE in the proxy environment." + ), + status_code=500, + ) + + class AzureAIAnthropicTokenCounter(BaseTokenCounter): """Token counter implementation for Azure AI Anthropic provider using the CountTokens API.""" @@ -65,18 +81,15 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): if not api_key: api_key = os.getenv("AZURE_AI_API_KEY") - # Get API base from deployment config or environment - api_base = litellm_params.get("api_base") - if not api_base: - api_base = os.getenv("AZURE_AI_API_BASE") - - if not api_key: - verbose_logger.warning("No Azure AI API key found for token counting") - return None + api_base: Final = ( + litellm_params.get("api_base") + or os.getenv("AZURE_AI_API_BASE") + or os.getenv("AZURE_API_BASE") + or litellm.api_base + ) if not api_base: - verbose_logger.warning("No Azure AI API base found for token counting") - return None + return _missing_api_base_response(request_model, model_to_use) try: result: Final = await azure_ai_anthropic_count_tokens_handler.handle_count_tokens_request( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index e96a2f20bbb..c73192fe50a 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -24,7 +24,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): def get_required_headers( self, - api_key: str, + api_key: str | None, litellm_params: dict[str, object] | None = None, ) -> dict[str, str]: """ @@ -45,12 +45,12 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): "Content-Type": "application/json", "anthropic-version": "2023-06-01", "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, - "x-api-key": api_key, # Azure AI Anthropic requires this header + **({"x-api-key": api_key} if api_key else {}), } # Also set up Azure auth headers for flexibility litellm_params = litellm_params or {} - if "api_key" not in litellm_params: + if api_key and "api_key" not in litellm_params: litellm_params["api_key"] = api_key litellm_params_obj: Final = GenericLiteLLMParams.model_validate(litellm_params) @@ -58,8 +58,9 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): # Get Azure auth headers (api-key or Authorization) azure_headers = BaseAzureLLM._base_validate_azure_environment(headers={}, litellm_params=litellm_params_obj) - # Merge Azure auth headers headers.update(azure_headers) + if "api-key" in azure_headers: + headers.setdefault("x-api-key", azure_headers["api-key"]) return headers diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 14f3884cc71..bdaea0b4104 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -2,20 +2,21 @@ OpenAI Token Counter implementation using the Responses API /input_tokens endpoint. """ -import os from collections.abc import Mapping, Sequence from typing import Any, Final +import openai from pydantic import JsonValue from typing_extensions import assert_never from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.llms.anthropic.experimental_pass_through.convertible_blocks import unconvertible_blocks_error from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.common_utils import OpenAIError, get_openai_credentials from litellm.llms.openai.responses.count_tokens.handler import ( OpenAICountTokensHandler, ) @@ -23,12 +24,25 @@ from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ResponsesCountRequestParts, ) +from litellm.llms.openai.workload_identity import ( + OpenAIWorkloadIdentityConfig, + get_workload_identity_bearer_token, + resolve_openai_workload_identity_config, +) from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests openai_count_tokens_handler: Final = OpenAICountTokensHandler() +def _workload_identity_config(litellm_params: Mapping[str, object]) -> OpenAIWorkloadIdentityConfig | None: + api_base: Final = litellm_params.get("api_base") + return resolve_openai_workload_identity_config( + api_key=None, + api_base=api_base if isinstance(api_base, str) else None, + ) + + class OpenAITokenCounter(BaseTokenCounter): """Token counter implementation for OpenAI provider using the Responses API.""" @@ -107,16 +121,27 @@ class OpenAITokenCounter(BaseTokenCounter): deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) - # Get OpenAI API key from deployment config or environment - api_key = litellm_params.get("api_key") - if not api_key: - api_key = os.getenv("OPENAI_API_KEY") + credentials: Final = get_openai_credentials( + api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key") + ) + credential: Final = credentials.api_key or _workload_identity_config(litellm_params) - if not api_key: - verbose_logger.warning("No OpenAI API key found for token counting") - return None + if credential is None: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=( + f"No OpenAI API key is set for model {request_model or model_to_use}, so tokens can't be " + "counted. Set api_key or workload identity on the deployment, or OPENAI_API_KEY in the proxy " + "environment." + ), + status_code=401, + ) - api_base: Final = litellm_params.get("api_base") + api_base: Final = credentials.api_base input_items, instructions, counted_tools, text = self._responses_request_parts( model_to_use=model_to_use, @@ -132,6 +157,11 @@ class OpenAITokenCounter(BaseTokenCounter): return None try: + api_key: Final = ( + credential + if isinstance(credential, str) + else await asyncify(get_workload_identity_bearer_token)(credential) + ) result: Final = await openai_count_tokens_handler.handle_count_tokens_request( model=model_to_use, input=input_items, @@ -161,6 +191,17 @@ class OpenAITokenCounter(BaseTokenCounter): error_message=e.message, status_code=e.status_code, ) + except openai.AuthenticationError as e: + verbose_logger.warning("OpenAI workload identity token exchange failed: %s", e.message) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=e.message, + status_code=401, + ) except Exception as e: verbose_logger.warning("Error calling OpenAI CountTokens API: %s", e) return TokenCountResponse( diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 2622446a699..210bd5a8bfd 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -29,6 +29,7 @@ from litellm.proxy.common_request_processing import ( create_response, log_llm_api_exception, proxy_exception_from_http_exception, + proxy_exception_from_route_error, resolve_litellm_call_id, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id @@ -336,10 +337,10 @@ async def count_tokens( messages: Final = data.get("messages", []) if not model_name: - raise HTTPException(status_code=400, detail={"error": "model parameter is required"}) + raise HTTPException(status_code=400, detail="model parameter is required") if not messages: - raise HTTPException(status_code=400, detail={"error": "messages parameter is required"}) + raise HTTPException(status_code=400, detail="messages parameter is required") # Create TokenCountRequest for the internal endpoint from litellm.proxy._types import TokenCountRequest @@ -357,13 +358,13 @@ async def count_tokens( messages_format="anthropic", output_format=_requested_output_format(data), ) - except HTTPException: - raise + except HTTPException as e: + return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) except ProxyException as e: return _anthropic_error_json_response(e, request) except Exception as e: log_llm_api_exception(e, litellm_call_id) - raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) + return _anthropic_error_json_response(proxy_exception_from_route_error(e), request) else: return {"input_tokens": token_response.total_tokens} diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6983786541c..1d642e72de5 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -609,6 +609,27 @@ def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, s ) +def proxy_exception_from_route_error(exc: Exception) -> ProxyException: + if isinstance(exc, ProxyException): + return exc + if isinstance(exc, HTTPException): + return proxy_exception_from_http_exception(exc, {}) + error_status: Final = error_status_code(exc, status.HTTP_500_INTERNAL_SERVER_ERROR) + message: Final = attribute_of(exc, "message", str(exc)) + carried_code: Final = attribute_of(exc, "code") + provider_fields: Final = attribute_of(exc, "provider_specific_fields") + return ProxyException( + message=redact_internal_details_from_client_message( + strip_bug_report_notice(message) if isinstance(message, str) else str(exc) + ), + type=openai_error_type(exc, error_status), + param=openai_error_param(exc), + code=error_status, + openai_code=carried_code if isinstance(carried_code, str) else None, + provider_specific_fields=provider_fields if isinstance(provider_fields, dict) else None, + ) + + def _collect_response_file_search_vector_store_ids(data: Mapping[str, object]) -> set[str]: vector_store_ids: Final[set[str]] = set() tools: Final = data.get("tools") diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 864ca2f9432..818fd1ad516 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,16 +1,22 @@ from collections.abc import Mapping from itertools import chain +from types import MappingProxyType from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response -from fastapi.responses import ORJSONResponse +from fastapi.responses import JSONResponse, ORJSONResponse from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types from litellm.proxy._types import * -from litellm.proxy._types import TokenCountRequest +from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + log_llm_api_exception, + proxy_exception_from_http_exception, + resolve_litellm_call_id, +) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.types.llms.vertex_ai import TokenCountDetailsResponse @@ -60,6 +66,37 @@ def _input_schema(declaration: _GoogleCountTokensFunctionDeclaration) -> object: return {"type": "object", "properties": {}} +_GOOGLE_RPC_STATUS: Final = MappingProxyType( + { + 400: "INVALID_ARGUMENT", + 401: "UNAUTHENTICATED", + 403: "PERMISSION_DENIED", + 404: "NOT_FOUND", + 409: "ABORTED", + 429: "RESOURCE_EXHAUSTED", + 499: "CANCELLED", + 501: "UNIMPLEMENTED", + 503: "UNAVAILABLE", + 504: "DEADLINE_EXCEEDED", + } +) + + +def _google_error_response(request: Request, error: ProxyException) -> JSONResponse: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does + ) + + status_code: Final = int(error.code) if error.code is not None and error.code.isdigit() else 500 + _close_dangling_otel_server_span(request, status_code, exc=error) + status: Final = _GOOGLE_RPC_STATUS.get(status_code, "INTERNAL" if status_code >= 500 else "FAILED_PRECONDITION") + return JSONResponse( + status_code=status_code, + content={"error": {"code": status_code, "message": error.message, "status": status}}, + headers=error.headers, + ) + + def _google_count_tokens_request(model_name: str, data: object) -> TokenCountRequest: from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter @@ -248,16 +285,26 @@ async def google_count_tokens(request: Request, model_name: str): data: Final = await _read_request_body(request=request) try: token_request: Final = _google_count_tokens_request(model_name, data) - except ValidationError as error: - raise HTTPException(status_code=400, detail="Invalid Google countTokens request body") from error + except ValidationError: + return _google_error_response( + request, + ProxyException( + message="Invalid Google countTokens request body", type="invalid_request_error", param=None, code=400 + ), + ) - # Call the internal token counter function with direct request flag set to False - token_response: Final = await internal_token_counter( - request=token_request, - call_endpoint=True, - ) - if token_response is None: - return TokenCountDetailsResponse(totalTokens=0, promptTokensDetails=[]) + try: + token_response: Final = await internal_token_counter(request=token_request, call_endpoint=True) + except ProxyException as error: + return _google_error_response(request, error) + except HTTPException as error: + return _google_error_response(request, proxy_exception_from_http_exception(error, {})) + except Exception as error: + log_llm_api_exception(error, resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))) + return _google_error_response( + request, + ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500), + ) original_response: Final[dict] = token_response.original_response or {} return TokenCountDetailsResponse( totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0, diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py new file mode 100644 index 00000000000..200f5e81673 --- /dev/null +++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py @@ -0,0 +1,205 @@ +import json +from collections.abc import Iterator +from typing import Final + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this synthetic request."}] +_PROVIDER_KEY_VARIABLES: Final = ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "AZURE_AI_API_KEY", + "AZURE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_TENANT_ID", + "OPENAI_API_KEY", +) + + +@pytest.fixture(scope="module") +def keyless_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-tokens-error-shapes"), + {}, + remove_environment=_PROVIDER_KEY_VARIABLES, + workers=1, + ) as owned, + ): + yield owned + + +@pytest.fixture(scope="module") +def azure_environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-tokens-azure-environment-key"), + {"AZURE_API_KEY": "synthetic-azure-env-key"}, + remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name != "AZURE_API_KEY"), + workers=1, + ) as owned, + ): + yield owned + + +def _anthropic(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client + ) + + +def _never_called(request: Request) -> Reply: + raise AssertionError(f"count reached the provider without a key: {request.target}") + + +@pytest.mark.parametrize( + ("deployment", "missing"), + ( + ({"model": "anthropic/claude-sonnet-4-6"}, "No Anthropic API key is set for model "), + ({"model": "azure_ai/claude-sonnet-4-6"}, "No Azure AI API key or Entra ID credential is set for model "), + ), + ids=("anthropic", "azure_ai_anthropic"), +) +def test_anthropic_count_without_a_provider_credential_is_an_anthropic_401( + keyless_proxy: OwnedProxy, deployment: dict[str, JsonValue], missing: str +) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(**deployment, api_base=wire.url, api_key=None, num_retries=0) + with pytest.raises(anthropic.AuthenticationError) as raised: + _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert raised.value.body == { + "type": "error", + "error": {"type": "authentication_error", "message": raised.value.body["error"]["message"]}, + } + assert missing in raised.value.message, raised.value.message + assert wire.drain() == () + + +def test_azure_ai_anthropic_count_sends_the_azure_environment_key_as_x_api_key( + azure_environment_key_proxy: OwnedProxy, +) -> None: + gateway: Final = azure_environment_key_proxy.gateway + + def count(request: Request) -> Reply: + assert request.headers.get("x-api-key") == "synthetic-azure-env-key", request.headers + return Reply(body=json.dumps({"input_tokens": 58}).encode()) + + with wire_server(count) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-sonnet-4-6", + api_base=wire.url, + api_key=None, + num_retries=0, + ) + result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert result.input_tokens == 58 + + +def test_azure_ai_anthropic_count_authenticates_with_an_entra_token(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + + def count(request: Request) -> Reply: + assert request.headers.get("authorization") == "Bearer synthetic-entra-token", request.headers + assert "x-api-key" not in request.headers, request.headers + return Reply(body=json.dumps({"input_tokens": 61}).encode()) + + with wire_server(count) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-sonnet-4-6", + api_base=wire.url, + api_key=None, + azure_ad_token="synthetic-entra-token", + num_retries=0, + ) + result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert result.input_tokens == 61 + assert len(wire.drain()) == 1 + + +def test_anthropic_count_request_errors_use_the_anthropic_envelope(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + upstream_error: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted reject"}} + with ( + wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", api_base=wire.url, api_key="synthetic-key", num_retries=0 + ) + missing_messages: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model}) + assert missing_messages.status_code == 400, missing_messages.text + assert _JSON_OBJECT.validate_json(missing_messages.content) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "messages parameter is required"}, + } + + rejected: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _MESSAGES} + ) + assert rejected.status_code == 400, rejected.text + body: Final = _JSON_OBJECT.validate_json(rejected.content) + assert set(body) <= {"type", "error", "request_id"}, body + assert body["error"] == upstream_error["error"], body + assert len(wire.drain()) == 1 + + +def test_openai_input_tokens_without_a_provider_key_is_an_openai_401(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key=None, num_retries=0 + ) + client: Final = openai.OpenAI( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client + ) + with pytest.raises(openai.AuthenticationError) as raised: + client.responses.input_tokens.count(model=model, input="Count this synthetic request.") + assert f"No OpenAI API key is set for model {model}" in raised.value.message + assert wire.drain() == () + + +def test_google_count_tokens_errors_use_the_google_envelope(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + upstream_error: Final = { + "error": {"code": 400, "message": "API key not valid. Please pass a valid API key.", "status": "INVALID_ARGUMENT"} + } + with ( + wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="gemini/gemini-2.5-flash", api_base=wire.url, api_key="synthetic-key", num_retries=0 + ) + rejected: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + {"contents": [{"role": "user", "parts": [{"text": "Count this synthetic request."}]}]}, + ) + assert rejected.status_code == 400, rejected.text + error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"] + assert isinstance(error, dict) and error["code"] == 400 and error["status"] == "INVALID_ARGUMENT", error + assert "API key not valid" in str(error["message"]), error + assert len(wire.drain()) == 1 + + malformed: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", {"contents": "not-a-list"}) + assert malformed.status_code == 400, malformed.text + assert _JSON_OBJECT.validate_json(malformed.content) == { + "error": {"code": 400, "message": "Invalid Google countTokens request body", "status": "INVALID_ARGUMENT"} + } diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index 9ecf681e039..1351680d617 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -426,8 +426,6 @@ class TestErrorLogCarriesCallId: @pytest.mark.asyncio async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture): - from fastapi import HTTPException - import litellm.proxy.anthropic_endpoints.endpoints as ep import litellm.proxy.proxy_server as proxy_server from litellm.proxy._types import UserAPIKeyAuth @@ -446,11 +444,10 @@ class TestErrorLogCarriesCallId: proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down")) ), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), - pytest.raises(HTTPException) as raised, ): - await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + response = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) - assert raised.value.status_code == 500 + assert response.status_code == 500 record = self._error_record(caplog) assert record.litellm_call_id == call_id assert call_id in record.getMessage() diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py index 451792109de..114bbb184c6 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -182,6 +182,29 @@ def test_google_count_tokens_unchanged(): assert body["totalTokens"] == 7 +def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope(): + try: + client = _build_test_client() + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + with patch( + "litellm.proxy.proxy_server.token_counter", + new_callable=AsyncMock, + side_effect=RuntimeError("tokenizer unavailable"), + ): + response = client.post( + "/v1beta/models/test-model:countTokens", + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + ) + + assert response.status_code == 500 + error = response.json()["error"] + assert error["code"] == 500 + assert error["status"] == "INTERNAL" + assert "tokenizer unavailable" not in error["message"] + + def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs(): """A provider that counts in Anthropic shape (Bedrock, Anthropic) carries no totalTokens key, so the normalized total must win over the raw response's missing field.""" try: diff --git a/tests/unit/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py index bcfc56577eb..169097ee47d 100644 --- a/tests/unit/llms/anthropic/test_count_tokens_oauth.py +++ b/tests/unit/llms/anthropic/test_count_tokens_oauth.py @@ -14,6 +14,12 @@ sys.path.insert( 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) ) +import httpx +import pytest +import respx + +import litellm +from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, ) @@ -84,3 +90,48 @@ class TestCountTokensOAuthHeaders: assert ( "oauth-2025-04-20" in beta_value ), f"oauth beta missing from OAuth headers: {beta_value}" + + +@pytest.mark.asyncio +async def test_a_global_anthropic_key_counts_like_chat_would_authenticate(monkeypatch): + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "anthropic_key", FAKE_REGULAR_KEY) + + with respx.mock(assert_all_called=True) as router: + route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 13}) + ) + result = await AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}}, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.total_tokens == 13 + assert route.calls.last.request.headers["x-api-key"] == FAKE_REGULAR_KEY + + +@pytest.mark.asyncio +async def test_no_anthropic_key_returns_a_401_count_error(monkeypatch): + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "anthropic_key", None) + monkeypatch.setattr(litellm, "api_key", None) + + result = await AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}}, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 401 + assert result.total_tokens == 0 diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py new file mode 100644 index 00000000000..86ad6c8d9d9 --- /dev/null +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -0,0 +1,212 @@ +import json +import operator +from typing import Final + +import httpx +import pytest +import respx +from pydantic import ValidationError + +import litellm +from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAnthropicTokenCounter +from litellm.types.utils import TokenCountResponse + +API_BASE: Final = "https://my-resource.services.ai.azure.com" +COUNT_URL: Final = f"{API_BASE}/anthropic/v1/messages/count_tokens" + + +async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + respx.post(COUNT_URL).mock(return_value=response) + return await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}}, + request_model="azure_ai/claude-sonnet-4-5", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "expected_total"), + [ + ({"input_tokens": 15}, 15), + ({"input_tokens": "15"}, 15), + ({"input_tokens": 15.0}, 15), + ({"input_tokens": True}, 1), + ({"input_tokens": -3}, -3), + ({}, 0), + ({"type": "error", "error": {"message": "429 rate limit"}}, 0), + ({"input_tokens": 7, "context_management": {"original_input_tokens": 9}, "nested": {"a": [1, None]}}, 7), + ], +) +async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload( + monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": expected_total, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "original_response": payload, + "error": False, + "error_message": None, + "status_code": None, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rejected_count", + [ + None, + 15.5, + "Request timed out", + "429 Too Many Requests", + [15], + [*range(403), "Request timed out"], + {"value": 429}, + ], +) +async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message( + monkeypatch: pytest.MonkeyPatch, rejected_count: object +) -> None: + with pytest.raises(ValidationError) as rejection: + TokenCountResponse.model_validate( + { + "total_tokens": rejected_count, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + } + ) + + result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count})) + + assert result is not None + assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)] + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "original_response": None, + "error": True, + "error_message": str(rejection.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False], +) +async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup( + monkeypatch: pytest.MonkeyPatch, payload: object +) -> None: + with pytest.raises(AttributeError) as failed_lookup: + operator.attrgetter("get")(payload) + + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "original_response": None, + "error": True, + "error_message": str(failed_lookup.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None: + assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("upstream_status", "upstream_text"), + [(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")], +) +async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text( + monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "original_response": None, + "error": True, + "error_message": upstream_text, + "status_code": upstream_status, + } + ) + + +@pytest.mark.asyncio +async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with pytest.raises(json.JSONDecodeError) as undecodable: + json.loads("not json") + + result: Final = await _count(monkeypatch, httpx.Response(200, text="not json")) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "original_response": None, + "error": True, + "error_message": f"CountTokens processing error: {undecodable.value}", + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("set_base", ("env", "global"), ids=("AZURE_API_BASE", "litellm.api_base")) +async def test_count_uses_the_shared_api_base_chat_would_use(monkeypatch: pytest.MonkeyPatch, set_base: str) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("AZURE_AI_API_BASE", raising=False) + monkeypatch.delenv("AZURE_API_BASE", raising=False) + if set_base == "env": + monkeypatch.setenv("AZURE_API_BASE", API_BASE) + else: + monkeypatch.setattr(litellm, "api_base", API_BASE) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key"}}, + request_model="azure_ai/claude-sonnet-4-5", + ) + assert result is not None + assert result.total_tokens == 6 + assert route.called diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py index 0f8d2aaba54..65a9f6f044d 100644 --- a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -388,3 +388,44 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_ "status_code": 500, } ) + + +@pytest.mark.asyncio +async def test_count_uses_the_global_openai_key_chat_would_use(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "openai_key", "sk-from-litellm-settings") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 9})) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.total_tokens == 9 + assert route.calls.last.request.headers["authorization"] == "Bearer sk-from-litellm-settings" + + +@pytest.mark.asyncio +async def test_count_sends_the_global_key_only_to_the_global_api_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "openai_key", "sk-private-gateway") + monkeypatch.setattr(litellm, "api_base", "https://gateway.internal.example/v1") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with respx.mock: + route: Final = respx.post("https://gateway.internal.example/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 4}) + ) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.total_tokens == 4 + assert route.calls.last.request.headers["authorization"] == "Bearer sk-private-gateway" diff --git a/tests/unit/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py index db107e00df0..aa0cad178f3 100644 --- a/tests/unit/llms/openai/test_openai_workload_identity.py +++ b/tests/unit/llms/openai/test_openai_workload_identity.py @@ -279,3 +279,56 @@ class TestResponsesValidateEnvironment: headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams() ) assert headers["Authorization"] == "Bearer None" + + +class TestCountTokens: + @pytest.mark.asyncio + @respx.mock + async def test_count_authenticates_with_the_workload_identity_bearer( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + mock_token_exchange() + count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 12}) + ) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.error is False + assert result.total_tokens == 12 + assert count_route.calls.last.request.headers["authorization"] == "Bearer exchanged-bearer-token" + + @pytest.mark.asyncio + @respx.mock + async def test_failed_mint_becomes_a_count_error_response( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + respx.post(TOKEN_EXCHANGE_URL).mock(return_value=httpx.Response(400, json={"error": "invalid_grant"})) + count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens") + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.error is True + assert result.status_code == 401 + assert result.total_tokens == 0 + assert not count_route.called diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index c1d437cd886..547cd795b24 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -327,22 +327,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 but non-Anthropic provider - response = await token_counter( - request=TokenCountRequest( - model="gpt-test", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) - - print("Non-Anthropic provider test response:", response) - - # Verify response structure - assert response.request_model == "gpt-test" - assert response.model_used == "gpt-4" - assert response.total_tokens > 0 - assert response.tokenizer_type == "openai_tokenizer" # Should use LiteLLM tokenizer + with pytest.raises(ProxyException) as missing_key: + await token_counter( + request=TokenCountRequest( + model="gpt-test", + messages=[{"role": "user", "content": "hello"}], + ), + call_endpoint=True, + ) + assert missing_key.value.code == "401" + assert "No OpenAI API key is set for model gpt-test" in missing_key.value.message @pytest.mark.asyncio @@ -352,12 +346,14 @@ async def test_anthropic_endpoint_error_handling(): """ from unittest.mock import MagicMock - from fastapi import HTTPException, Request + from fastapi import Request from litellm.proxy.anthropic_endpoints.endpoints import count_tokens # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} + mock_request.scope = {} mock_user_api_key_dict = MagicMock() # Test missing model parameter @@ -375,12 +371,13 @@ async def test_anthropic_endpoint_error_handling(): anthropic_endpoints._read_request_body = mock_read_request_body try: - # Should raise HTTPException for missing model - with pytest.raises(HTTPException) as exc_info: - await count_tokens(mock_request, mock_user_api_key_dict) + response = await count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 400 - assert "model parameter is required" in str(exc_info.value.detail) + assert response.status_code == 400 + assert json.loads(response.body) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "model parameter is required"}, + } print("✅ Error handling test passed!") @@ -471,7 +468,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): mock_router.model_list = [ { "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, + "litellm_params": {"model": "databricks/gpt-4"}, "model_info": {}, } ] @@ -480,7 +477,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): mock_router.async_get_available_deployment = AsyncMock( return_value={ "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, + "litellm_params": {"model": "databricks/gpt-4"}, "model_info": {}, } ) From b1037d2121115c6e869c7e7fab935104781716e6 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:46:48 -0700 Subject: [PATCH 18/36] fix(proxy)!: count prompt-only token_counter requests with the provider when call_endpoint is set (#46839) Backport of #46839 to stable/1.104.x. Cherry-picked from 2dc6aafa6baf8dabdac58f5ba70c384b8dcb3e09 (main). --- litellm/proxy/proxy_server.py | 34 ++- .../test_vertex_gemini_count_tokens_wire.py | 114 ++++++- .../proxy/proxy_server/test_routes_utils.py | 7 +- tests/unit/proxy/test_proxy_token_counter.py | 279 +++++++++++++++++- 4 files changed, 421 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f5810a96a9c..ee1e8647bf2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13657,6 +13657,16 @@ async def _try_provider_token_count( return result +_PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( + {litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value} +) + + +def _drops_params(deployment: Mapping[str, object] | None) -> bool: + litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) + return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) + + def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -13723,8 +13733,8 @@ async def count_request_tokens( litellm_model_name = litellm_model_name.split("/", 1)[1] model_to_use: str = ( - litellm_model_name or request.model - ) # use litellm model name, if it's not avalable then fallback to request.model + litellm_model_name if isinstance(litellm_model_name, str) and litellm_model_name else request.model + ) # Try provider-specific token counting first - only for non-direct requests (from provider endpoints) provider_counter: BaseTokenCounter | None = None @@ -13735,13 +13745,23 @@ async def count_request_tokens( if _model is not None: model_to_use = _model + conversation_empty: Final = not messages and not contents + if prompt and not conversation_empty and not _drops_params(deployment): + raise ProxyException( + message="prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt", + type="invalid_request_error", + param="prompt", + code=400, + ) + prompt_only: Final = bool(prompt) and conversation_empty if provider_counter is not None: + prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES result: Final = await _try_provider_token_count( provider_counter=provider_counter, custom_llm_provider=custom_llm_provider, model_to_use=model_to_use, - messages=messages, - contents=contents, + messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages, + contents=None if prompt_as_message else contents, deployment=deployment, request_model=request.model, tools=tools, @@ -13774,8 +13794,10 @@ async def count_request_tokens( tokenizer_used: Final = str(_tokenizer_used["type"]) system_message: Final = _system_message(system) + prompt_turn: Final = 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, messages + Sequence[AllMessageValues] | None, + [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else messages), ) counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages) @@ -13785,7 +13807,7 @@ async def count_request_tokens( ) total_tokens: Final = await offload_token_count(litellm.token_counter)( model=model_to_use, - text=prompt, + text=prompt if counted_messages is None else None, messages=counted_messages, tools=counted_tools, custom_tokenizer=_tokenizer_used, diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 5fe82d37e9e..8de93256f6c 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -65,7 +65,7 @@ class _GeminiPeer: return requests -def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str: +def _deployment(gateway: Gateway, scenario: Scenario, api_base: str, **parameters: JsonValue) -> str: return scenario.model( model=f"vertex_ai/{_MODEL}", api_base=api_base, @@ -73,6 +73,7 @@ def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str: vertex_project=_PROJECT, vertex_location=_LOCATION, vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + **parameters, ) @@ -220,8 +221,8 @@ _CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingPro ) -def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str: - return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key") +def _ai_studio_deployment(scenario: Scenario, api_base: str, **parameters: JsonValue) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key", **parameters) def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: @@ -594,3 +595,110 @@ def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gatew tuple(_json_body(request) for request in requests) == ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9 ) + + +_CONTENTS: Final[JsonValue] = [{"role": "user", "parts": [{"text": _USER}]}] +_OPENAI_TOOLS: Final[JsonValue] = [ + { + "type": "function", + "function": {"name": "noop", "description": "Return no output", "parameters": {"type": "object"}}, + } +] + + +def _provider_deployment( + gateway: Gateway, scenario: Scenario, provider: str, api_base: str, **parameters: JsonValue +) -> str: + if provider == "vertex_ai": + return _deployment(gateway, scenario, api_base, **parameters) + return _ai_studio_deployment(scenario, api_base, **parameters) + + +def _count(gateway: Gateway, body: Mapping[str, JsonValue], call_endpoint: bool = True) -> httpx.Response: + return gateway.request( + "POST", "/utils/token_counter", body, params={"call_endpoint": "true" if call_endpoint else "false"} + ) + + +def _total_tokens(response: httpx.Response) -> JsonValue: + return _JSON_OBJECT.validate_json(response.content)["total_tokens"] + + +@pytest.mark.parametrize( + ("provider", "empty_conversation"), + (("vertex_ai", {}), ("gemini", {}), ("gemini", {"messages": []}), ("gemini", {"contents": []})), + ids=("vertex-prompt-only", "ai-studio-prompt-only", "ai-studio-empty-messages", "ai-studio-empty-contents"), +) +def test_utils_token_counter_sends_a_prompt_as_the_user_turn_of_its_messages_form( + gateway: Gateway, provider: str, empty_conversation: Mapping[str, JsonValue] +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _provider_deployment(gateway, scenario, provider, wire.url) + prompted: Final = _count(gateway, {"model": model, "prompt": _USER, **empty_conversation}) + prompted_requests: Final = _count_requests(wire) + as_message: Final = _count(gateway, {"model": model, "messages": _MESSAGES}) + message_requests: Final = _count_requests(wire) + assert prompted.status_code == 200, prompted.text + assert as_message.status_code == 200, as_message.text + assert _total_tokens(prompted) == _COUNT, prompted.text + assert len(prompted_requests) == 1 + assert len(message_requests) == 1 + assert _counted_contents(_json_body(prompted_requests[0])) == _counted_contents(_json_body(message_requests[0])) + + +@pytest.mark.parametrize( + ("conversation", "call_endpoint"), + (({"messages": _MESSAGES}, True), ({"contents": _CONTENTS}, True), ({"messages": _MESSAGES}, False)), + ids=("provider-count-messages", "provider-count-contents", "local-count-messages"), +) +def test_utils_token_counter_rejects_a_prompt_next_to_a_conversation_without_calling_the_provider( + gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url) + response: Final = _count(gateway, {"model": model, "prompt": "hi", **conversation}, call_endpoint) + requests: Final = _count_requests(wire) + assert response.status_code == 400, response.text + assert _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_json(response.content)["error"])["param"] == "prompt" + assert requests == () + + +@pytest.mark.parametrize("call_endpoint", (True, False), ids=("provider-count", "local-count")) +def test_utils_token_counter_counts_only_the_conversation_when_drop_params_drops_the_prompt( + gateway: Gateway, call_endpoint: bool +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=True) + mixed: Final = _count(gateway, {"model": model, "prompt": "hi", "messages": _MESSAGES}, call_endpoint) + mixed_requests: Final = _count_requests(wire) + alone: Final = _count(gateway, {"model": model, "messages": _MESSAGES}, call_endpoint) + alone_requests: Final = _count_requests(wire) + assert mixed.status_code == 200, mixed.text + assert alone.status_code == 200, alone.text + assert _total_tokens(mixed) == _total_tokens(alone) + assert [_counted_contents(_json_body(request)) for request in mixed_requests] == [ + _counted_contents(_json_body(request)) for request in alone_requests + ] + + +@pytest.mark.parametrize( + "extra", + ({"system": "Follow the system instruction."}, {"tools": _OPENAI_TOOLS}), + ids=("system", "tools"), +) +def test_utils_token_counter_counts_a_local_prompt_with_system_or_tools_like_its_messages_form( + gateway: Gateway, extra: Mapping[str, JsonValue] +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url) + prompted: Final = _count(gateway, {"model": model, "prompt": _USER, **extra}, call_endpoint=False) + as_message: Final = _count(gateway, {"model": model, "messages": _MESSAGES, **extra}, call_endpoint=False) + requests: Final = _count_requests(wire) + assert prompted.status_code == 200, prompted.text + assert as_message.status_code == 200, as_message.text + assert _total_tokens(prompted) == _total_tokens(as_message) + assert requests == () diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py index 6cd613197ee..fc5ab9475eb 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py @@ -382,8 +382,7 @@ def test_token_counter_fallback_counts_tools_system_and_anthropic_blocks(client, assert full > bare -def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, monkeypatch): - """Regression: a ``prompt`` request carrying ``tools`` but no ``messages`` still counts, because the fallback attaches tools only when counting messages (``token_counter`` rejects tools on the text path).""" +def test_token_counter_fallback_counts_a_prompt_with_tools_like_its_messages_form(client, auth_as, monkeypatch): monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False) prompt = "count the tokens in this sentence please" @@ -404,4 +403,6 @@ def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, ) assert response.status_code == 200, response.text - assert response.json()["total_tokens"] == litellm.token_counter(model="claude-fable-5", text=prompt) + assert response.json()["total_tokens"] == litellm.token_counter( + model="claude-fable-5", messages=[{"role": "user", "content": prompt}], tools=tools + ) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 547cd795b24..869e7eb60d2 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -9,7 +9,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from dotenv import load_dotenv -from pydantic import JsonValue +from pydantic import BaseModel, ConfigDict, JsonValue load_dotenv() @@ -21,10 +21,13 @@ from fastapi.responses import JSONResponse import litellm from litellm import Router +from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler +from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.anthropic_endpoints.endpoints import ( count_tokens as anthropic_count_tokens, @@ -1174,3 +1177,277 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): finally: anthropic_endpoints._read_request_body = original_read_request_body proxy_server.count_request_tokens = original_token_counter + + +class _GeminiPart(BaseModel): + model_config = ConfigDict(frozen=True) + text: str = "" + + +class _GeminiContent(BaseModel): + model_config = ConfigDict(frozen=True) + parts: tuple[_GeminiPart, ...] = () + + +class _GeminiCountTokensBody(BaseModel): + model_config = ConfigDict(frozen=True) + contents: tuple[_GeminiContent, ...] = () + + +def _word_count(content: _GeminiContent) -> int: + return sum(len(part.text.split()) for part in content.parts) + + +def _gemini_word_counting_peer(request: httpx.Request) -> httpx.Response: + body: Final = _GeminiCountTokensBody.model_validate_json(request.content) + total: Final = sum(_word_count(content) for content in body.contents) + return httpx.Response(200, json={"totalTokens": total}, request=request) + + +def _unavailable_peer(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"error": {"message": "unavailable"}}, request=request) + + +def _gemini_router(drop_params: bool = False) -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-counted", + "litellm_params": { + "model": "gemini/gemini-3.8-flash", + "api_key": "test-key", + "api_base": "https://count-tokens.test", + "drop_params": drop_params, + }, + } + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "text", ["hello there", "The quick brown fox jumps over the lazy dog. " * 50], ids=["short", "long"] +) +async def test_call_endpoint_counts_a_prompt_like_the_same_text_sent_as_a_user_message( + text: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=text), call_endpoint=True + ) + from_message: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": text}]), + call_endpoint=True, + ) + + assert from_prompt.error is False + assert from_prompt.total_tokens == from_message.total_tokens == len(text.split()) + + +class _PartnerMessage(BaseModel): + model_config = ConfigDict(frozen=True) + content: str = "" + + +class _PartnerCountRequest(BaseModel): + model_config = ConfigDict(frozen=True) + messages: tuple[_PartnerMessage, ...] = () + + +async def _vertex_partner_word_count( + model: str, request_data: JsonValue, litellm_params: JsonValue +) -> dict[str, JsonValue]: + body: Final = _PartnerCountRequest.model_validate(request_data) + return {"input_tokens": sum(len(message.content.split()) for message in body.messages)} + + +@pytest.mark.asyncio +async def test_call_endpoint_counts_a_prompt_on_a_vertex_ai_deployment(monkeypatch: pytest.MonkeyPatch) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "vertex-counted", + "litellm_params": { + "model": "vertex_ai/claude-sonnet-5-5", + "vertex_project": "test-project", + "vertex_location": "us-east5", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + text: Final = "The quick brown fox jumps over the lazy dog." + + with patch.object( + VertexAIPartnerModelsTokenCounter, "handle_count_tokens_request", side_effect=_vertex_partner_word_count + ): + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="vertex-counted", prompt=text), call_endpoint=True + ) + + assert (from_prompt.error, from_prompt.total_tokens) == (False, 9) + + +_FOX: Final = "The quick brown fox jumps over the lazy dog." + + +_CONVERSATIONS: Final = ( + {"messages": [{"role": "user", "content": _FOX}]}, + {"contents": [{"role": "user", "parts": [{"text": _FOX}]}]}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("conversation", _CONVERSATIONS, ids=["messages", "contents"]) +@pytest.mark.parametrize("call_endpoint", [True, False], ids=["provider", "local"]) +async def test_prompt_next_to_a_conversation_answers_400_without_drop_params( + conversation: dict[str, JsonValue], call_endpoint: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with ( + patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), + pytest.raises(ProxyException) as rejected, + ): + await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint + ) + + assert (rejected.value.code, rejected.value.param) == ("400", "prompt") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("conversation", "call_endpoint", "global_drop_params", "deployment_drop_params"), + [ + (_CONVERSATIONS[0], True, False, True), + (_CONVERSATIONS[1], True, False, True), + (_CONVERSATIONS[0], False, False, True), + (_CONVERSATIONS[0], False, True, False), + ], + ids=[ + "provider_messages", + "provider_contents", + "local_messages", + "local_messages_global_setting", + ], +) +async def test_drop_params_counts_a_conversation_without_the_prompt( + conversation: dict[str, JsonValue], + call_endpoint: bool, + global_drop_params: bool, + deployment_drop_params: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router(drop_params=deployment_drop_params)) + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + with_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint + ) + conversation_only: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", **conversation), call_endpoint=call_endpoint + ) + + assert with_prompt.error is False + assert with_prompt.total_tokens == conversation_only.total_tokens > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("empty_conversation", [{"messages": []}, {"contents": []}], ids=["messages", "contents"]) +async def test_call_endpoint_prompt_next_to_an_empty_conversation_counts_the_prompt( + empty_conversation: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + counted: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation), call_endpoint=True + ) + + assert (counted.error, counted.total_tokens) == (False, 9) + + +@pytest.mark.asyncio +async def test_local_count_of_a_prompt_next_to_empty_messages_counts_the_prompt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + + with_empty: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX, messages=[])) + alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) + + assert with_empty.error is False + assert with_empty.total_tokens == alone.total_tokens > 0 + + +_WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chat_fields", + [{"system": "You answer in one short sentence. " * 20}, {"tools": [_WEATHER_TOOL]}], + ids=["system", "tools"], +) +@pytest.mark.parametrize("empty_conversation", [{}, {"messages": []}], ids=["no_messages", "empty_messages"]) +async def test_local_count_of_a_prompt_with_system_or_tools_counts_them_like_the_message_form( + chat_fields: dict[str, JsonValue], empty_conversation: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation, **chat_fields) + ) + from_message: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": _FOX}], **chat_fields) + ) + prompt_alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) + + assert from_prompt.error is False + assert from_prompt.total_tokens == from_message.total_tokens > prompt_alone.total_tokens + + +@pytest.mark.asyncio +async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_needs_messages( + 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.") + + 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) + 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, + ) From a4be1c1d137afc1c15cbc720cfb123927f4c4854 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:48:20 -0700 Subject: [PATCH 19/36] ci(cost map): classify token_count_model as an object key in the schema generator (#47057) Backport of #47057 to stable/1.104.x. Cherry-picked from 6e191a58f1780e4d9e267f47e65208742c713b17 (main). --- ci_cd/generate_model_prices_schema.py | 16 +++++++++++----- tests/unit/test_model_prices_schema.py | 25 +++++++++++++++++++------ 2 files changed, 30 insertions(+), 11 deletions(-) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index e429aae36b9..359df8e6e32 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -19,6 +19,7 @@ NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0} NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0} BOOLEAN: JsonSchema = {"type": "boolean"} STRING: JsonSchema = {"type": "string"} +NON_BLANK_STRING: Final[JsonSchema] = {"type": "string", "pattern": r"\S"} TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"} WEEKDAY_PATTERN: Final = ( r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$" @@ -103,6 +104,16 @@ OBJECT_KEYS: dict[str, JsonSchema] = { "type": "object", "description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).", }, + "token_count_model": { + "type": "object", + "required": ["model"], + "additionalProperties": False, + "properties": { + "model": NON_BLANK_STRING, + "region": NON_BLANK_STRING, + }, + "description": "Model that counts tokens for this one, and the AWS region its count API runs in.", + }, } ARRAY_KEYS: dict[str, JsonSchema] = { @@ -287,11 +298,6 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]: "enum": ["none", "minimal", "low", "medium", "high", "xhigh"], }, "comment": STRING, - "token_count_fallback_model": { - "type": "string", - "minLength": 1, - "description": "Model that counts tokens for this one when the provider cannot count it.", - }, "audio_transcription_config": STRING, "vertex_ai_audio_api": { "type": "string", diff --git a/tests/unit/test_model_prices_schema.py b/tests/unit/test_model_prices_schema.py index ccc379ffe4f..211a43d2574 100644 --- a/tests/unit/test_model_prices_schema.py +++ b/tests/unit/test_model_prices_schema.py @@ -181,16 +181,23 @@ TOKEN_COUNT_ENTRY: Final = MappingProxyType( "litellm_provider": "bedrock_converse", "mode": "chat", "token_count_framing": 78, - "token_count_fallback_model": "us.anthropic.claude-opus-5-5", + "token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"}, } ) -@pytest.mark.parametrize("framing", [0, 78]) -def test_generator_classifies_token_count_settings(framing: int): +@pytest.mark.parametrize( + "setting", + [ + {"token_count_framing": 0}, + {"token_count_framing": 78}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5"}}, + ], +) +def test_generator_classifies_token_count_settings(setting: dict): generator = load_generator() schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) - assert build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, "token_count_framing": framing}}) + assert build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}}) @pytest.mark.parametrize( @@ -199,8 +206,14 @@ def test_generator_classifies_token_count_settings(framing: int): {"token_count_framing": -1}, {"token_count_framing": "78"}, {"token_count_framing": 1.5}, - {"token_count_fallback_model": ""}, - {"token_count_fallback_model": 5}, + {"token_count_model": {"model": ""}}, + {"token_count_model": {"model": " "}}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": ""}}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": " "}}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": 5}}, + {"token_count_model": {"region": "us-east-1"}}, + {"token_count_model": "us.anthropic.claude-opus-5-5"}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "zone": "us-east-1"}}, ], ) def test_generated_token_count_schema_rejects_malformed_settings(setting: dict): From da5a9ffd2e25e7eecad646c1bc9771842331106a Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:49:48 -0700 Subject: [PATCH 20/36] ci(cost map): remove the generated model prices JSON schema (#47064) Backport of #47064 to stable/1.104.x. Cherry-picked from ae6f7f661355d0764031678d014c29f765897b68 (main). Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auto_update_price_and_context_window.yml | 5 +- .github/workflows/sync-together-ai-models.yml | 7 +- ci_cd/cost_map_guard.py | 45 +- ci_cd/generate_model_prices_schema.py | 424 ------------------ litellm-rust/crates/model-catalog/AGENTS.md | 1 - tests/unit/test_cost_map_guard.py | 52 +-- ...> test_model_prices_and_context_window.py} | 198 -------- 7 files changed, 27 insertions(+), 705 deletions(-) delete mode 100644 ci_cd/generate_model_prices_schema.py rename tests/unit/{test_model_prices_schema.py => test_model_prices_and_context_window.py} (67%) diff --git a/.github/workflows/auto_update_price_and_context_window.yml b/.github/workflows/auto_update_price_and_context_window.yml index 7e40a860ee9..838ba35f831 100644 --- a/.github/workflows/auto_update_price_and_context_window.yml +++ b/.github/workflows/auto_update_price_and_context_window.yml @@ -24,12 +24,9 @@ jobs: - name: Update JSON Data run: | uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py" - - name: Regenerate JSON Schema - run: | - uv run --frozen python ci_cd/generate_model_prices_schema.py - name: Create Pull Request run: | - git add model_prices_and_context_window.json model_prices_and_context_window.schema.json + git add model_prices_and_context_window.json git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')" gh pr create --title "Update model_prices_and_context_window.json file" \ --body "Automated update for model_prices_and_context_window.json" \ diff --git a/.github/workflows/sync-together-ai-models.yml b/.github/workflows/sync-together-ai-models.yml index 18d5e3de1eb..5325510e45e 100644 --- a/.github/workflows/sync-together-ai-models.yml +++ b/.github/workflows/sync-together-ai-models.yml @@ -41,10 +41,6 @@ jobs: uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md" env: TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }} - - name: Regenerate the JSON schema - if: steps.existing.outputs.open_pr == '' - run: | - uv run --frozen python ci_cd/generate_model_prices_schema.py - name: Create a pull request when the registry changed if: steps.existing.outputs.open_pr == '' run: | @@ -57,8 +53,7 @@ jobs: git config user.email "41898282+github-actions[bot]@users.noreply.github.com" git checkout -b "$branch" git add model_prices_and_context_window.json \ - litellm/model_prices_and_context_window_backup.json \ - model_prices_and_context_window.schema.json + litellm/model_prices_and_context_window_backup.json git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')" gh auth setup-git git push origin "$branch" diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py index 14eecf01461..8e34e79a354 100644 --- a/ci_cd/cost_map_guard.py +++ b/ci_cd/cost_map_guard.py @@ -1,11 +1,11 @@ """Guard the cost map on pull requests. -Every pull request whose diff against its merge base touches one of the three cost map files gets the file -checks: the files parse, the backup copy matches the root file, and the JSON schema is in sync and validates the -map. A pull request that leaves all three untouched skips them, since merging it keeps the base branch's copies -and its head tree only carries whatever state the branch was cut from. Pull requests from the cost map sync bot -(branches named litellm_cost_map_sync_*) always get the file checks and additionally may only touch those three -files and may only add or update models. +Every pull request whose diff against its merge base touches one of the two cost map files gets the file +checks: the files parse and the backup copy matches the root file. A pull request that leaves both untouched +skips them, since merging it keeps the base branch's copies and its head tree only carries whatever state the +branch was cut from. Pull requests from the cost map sync bot (branches named litellm_cost_map_sync_*) always +get the file checks and additionally may only touch those two files or the legacy schema +path and may only add or update models. """ from __future__ import annotations @@ -19,13 +19,12 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final -from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors - COST_MAP_PATH: Final = "model_prices_and_context_window.json" BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json" -SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json" -GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH) +LEGACY_SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json" +GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH) BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_" +SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"}) EVALUATION_OUTPUT_PRICE_SOURCES: Final[Mapping[str, str]] = MappingProxyType({}) CostMap = dict[str, object] @@ -35,7 +34,6 @@ CostMap = dict[str, object] class Snapshot: cost_map: str backup: str - schema: str def _parse_object(text: str, path: str) -> CostMap | str: @@ -46,13 +44,6 @@ def _parse_object(text: str, path: str) -> CostMap | str: return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root" -def _rendered_schema(cost_map: CostMap) -> str: - try: - return render(build_schema(cost_map)) - except SystemExit as error: - return str(error) - - def _evaluation_output_price_failure( key: str, entry: object, evaluation_output_price_sources: Mapping[str, str] ) -> str | None: @@ -78,29 +69,13 @@ def _file_failures( *, evaluation_output_price_sources: Mapping[str, str] = EVALUATION_OUTPUT_PRICE_SOURCES, ) -> tuple[str, ...]: - schema_text: Final = _rendered_schema(head_map) - if not schema_text.startswith("{"): - return (schema_text,) backup_failure: Final = ( () if head.backup == head.cost_map else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",) ) - schema_failure: Final = ( - () - if head.schema == schema_text - else ( - f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; " - "run `python ci_cd/generate_model_prices_schema.py` and commit the result", - ) - ) return ( *backup_failure, - *schema_failure, - *( - f"{COST_MAP_PATH} does not validate against its schema: {error}" - for error in validation_errors(head_map, json.loads(schema_text))[:20] - ), *( failure for key, entry in head_map.items() @@ -135,7 +110,7 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str *( f"bot PRs may only change the cost map files, not {path}" for path in changed_files - if path not in GUARDED_PATHS + if path not in GUARDED_PATHS and path != LEGACY_SCHEMA_PATH ), *(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map), *(f"bot PRs may not remove fields: {ref}" for ref in removed_fields), diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py deleted file mode 100644 index 359df8e6e32..00000000000 --- a/ci_cd/generate_model_prices_schema.py +++ /dev/null @@ -1,424 +0,0 @@ -from __future__ import annotations - -import json -import sys -from pathlib import Path -from typing import Final, Optional - -import jsonschema - -REPO_ROOT = Path(__file__).parent.parent -PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json" -SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json" - -SPECIAL_ROOT_KEYS = frozenset({"sample_spec", "fallback_generalizations"}) - -JsonSchema = dict - -NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0} -NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0} -BOOLEAN: JsonSchema = {"type": "boolean"} -STRING: JsonSchema = {"type": "string"} -NON_BLANK_STRING: Final[JsonSchema] = {"type": "string", "pattern": r"\S"} -TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"} -WEEKDAY_PATTERN: Final = ( - r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$" -) - -EXTRA_BOOLEAN_KEYS = frozenset( - { - "gemini_native_audio", - "gemini_audio_only_live", - "uses_embed_content", - "use_openai_responses_path", - "bedrock_converse_supports_strict_tools", - "thinking_always_on", - } -) - -HOURS_UTC: Final[JsonSchema] = { - "description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.', - "oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}], -} - -OFF_PEAK_WINDOW: Final[JsonSchema] = { - "type": "object", - "properties": { - "hours_utc": HOURS_UTC, - "weekdays": { - "type": "array", - "description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.", - "items": { - "oneOf": [ - {"type": "integer", "minimum": 1, "maximum": 7}, - {"type": "string", "pattern": WEEKDAY_PATTERN}, - ] - }, - "minItems": 1, - }, - }, - "required": ["hours_utc"], - "additionalProperties": False, -} - -OBJECT_KEYS: dict[str, JsonSchema] = { - "off_peak_pricing": { - "type": "object", - "description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.", - "properties": { - "hours_utc": HOURS_UTC, - "windows": {"type": "array", "items": OFF_PEAK_WINDOW, "minItems": 1}, - "weekday_timezone": { - "type": "string", - "description": "IANA zone the weekdays of each window are read on; defaults to UTC.", - }, - "input_cost_per_token": NONNEG_NUMBER, - "output_cost_per_token": NONNEG_NUMBER, - "output_cost_per_reasoning_token": NONNEG_NUMBER, - "cache_read_input_token_cost": NONNEG_NUMBER, - "cache_creation_input_token_cost": NONNEG_NUMBER, - }, - "anyOf": [{"required": ["hours_utc"]}, {"required": ["windows"]}], - "additionalProperties": False, - }, - "search_context_cost_per_query": { - "type": "object", - "description": "USD cost per web search query, keyed by search context size.", - "properties": { - "search_context_size_low": NONNEG_NUMBER, - "search_context_size_medium": NONNEG_NUMBER, - "search_context_size_high": NONNEG_NUMBER, - }, - "additionalProperties": False, - }, - "guardrail_cost_per_unit": { - "type": "object", - "description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).", - "additionalProperties": NONNEG_NUMBER, - }, - "metadata": { - "type": "object", - "description": "Free-form notes about the entry (e.g. pricing derivation).", - }, - "provider_specific_entry": { - "type": "object", - "description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).", - }, - "token_count_model": { - "type": "object", - "required": ["model"], - "additionalProperties": False, - "properties": { - "model": NON_BLANK_STRING, - "region": NON_BLANK_STRING, - }, - "description": "Model that counts tokens for this one, and the AWS region its count API runs in.", - }, -} - -ARRAY_KEYS: dict[str, JsonSchema] = { - "supported_audio_formats": { - "type": "array", - "description": "Audio container formats the model can return.", - "items": {"type": "string", "enum": ["mp3", "wav"]}, - }, - "supported_endpoints": { - "type": "array", - "description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.", - "items": STRING, - }, - "supported_modalities": { - "type": "array", - "description": "Input modalities the model accepts.", - "items": {"type": "string", "enum": ["text", "image", "audio", "video"]}, - }, - "supported_output_modalities": { - "type": "array", - "description": "Output modalities the model can produce.", - "items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]}, - }, - "reasoning_effort_levels": { - "type": "array", - "description": "Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.", - "items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]}, - }, - "supported_regions": { - "type": "array", - "description": "Cloud regions the model is available in ('global' or region ids).", - "items": STRING, - }, - "tiered_pricing": { - "type": "array", - "description": "Context-length or result-count tiered rates; each tier's costs apply within its range.", - "items": { - "type": "object", - "properties": { - "range": { - "type": "array", - "description": "[min, max] prompt-token span this tier applies to.", - "items": NONNEG_NUMBER, - "minItems": 2, - "maxItems": 2, - }, - "max_results_range": { - "type": "array", - "description": "[min, max] result-count span this tier applies to (search models).", - "items": NONNEG_NUMBER, - "minItems": 2, - "maxItems": 2, - }, - "input_cost_per_token": NONNEG_NUMBER, - "output_cost_per_token": NONNEG_NUMBER, - "output_cost_per_reasoning_token": NONNEG_NUMBER, - "cache_read_input_token_cost": NONNEG_NUMBER, - "cache_creation_input_token_cost": NONNEG_NUMBER, - "input_cost_per_query": NONNEG_NUMBER, - }, - "additionalProperties": False, - }, - }, -} - -INTEGER_KEYS: dict[str, JsonSchema] = { - "max_tokens": { - **NONNEG_INTEGER, - "description": "Legacy field: max output tokens if the provider specifies it, else max input tokens.", - }, - "max_input_tokens": { - **NONNEG_INTEGER, - "description": "Maximum prompt/context tokens the model accepts.", - }, - "max_output_tokens": { - **NONNEG_INTEGER, - "description": "Maximum tokens the model can generate in one response.", - }, - "output_vector_size": { - **NONNEG_INTEGER, - "description": "Embedding dimension for embedding models.", - }, - "prompt_cache_min_tokens": { - **NONNEG_INTEGER, - "description": "Smallest prefix the provider will actually cache; absent means the provider default applies.", - }, - "token_count_framing": { - **NONNEG_INTEGER, - "description": "Tokens added once per request to a local count estimate, for the provider's chat framing.", - }, - "tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."}, - "rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."}, -} - -NUMBER_KEYS: dict[str, JsonSchema] = { - "regional_processing_uplift_multiplier_eu": { - "type": "number", - "minimum": 1, - "description": "Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%).", - }, - "regional_processing_uplift_multiplier_us": { - "type": "number", - "minimum": 1, - "description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).", - }, - "regional_endpoint_uplift_multiplier": { - "type": "number", - "minimum": 1, - "description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%).", - }, -} - -COST_DESCRIPTIONS: dict[str, str] = { - "input_cost_per_token": "USD per prompt token.", - "output_cost_per_token": "USD per generated token.", - "output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.", - "google_maps_grounding_cost_per_query": ( - "USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit." - ), - "cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.", - "cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.", - "input_cost_per_token_batches": "USD per prompt token via the provider's batch API.", - "output_cost_per_token_batches": "USD per generated token via the provider's batch API.", -} - - -def cost_description(key: str) -> Optional[str]: - if key in COST_DESCRIPTIONS: - return COST_DESCRIPTIONS[key] - if key.endswith("_flex"): - return "Flex service-tier rate for the same-named base field." - if key.endswith("_priority"): - return "Priority service-tier rate for the same-named base field." - if "_above_" in key: - return "Rate applied once the prompt exceeds the token threshold in the field name." - return None - - -def cost_schema(key: str) -> JsonSchema: - description = cost_description(key) - return {**NONNEG_NUMBER, "description": description} if description else dict(NONNEG_NUMBER) - - -def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]: - return { - "litellm_provider": { - "type": "string", - "description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.", - }, - "mode": { - "type": "string", - "description": "Primary API surface / task type of the model.", - "enum": list(modes), - }, - "source": { - "type": "string", - "description": "URL of the provider pricing/model page this entry was taken from.", - }, - "deprecation_date": { - "type": "string", - "description": "Date the provider deprecates the model, YYYY-MM-DD.", - "format": "date", - "pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$", - }, - "web_search_billing_unit": { - "type": "string", - "description": "Whether web search is billed per query or per prompt.", - "enum": ["per_query", "per_prompt"], - }, - "bedrock_output_config_effort_ceiling": { - "type": "string", - "description": "Highest reasoning effort the Bedrock output_config accepts for this model.", - "enum": ["low", "medium", "high", "max", "xhigh"], - }, - "default_reasoning_effort": { - "type": "string", - "description": ( - "Reasoning effort the provider applies when the request omits reasoning_effort. " - "Gates whether a non-default temperature or the top_p/logprobs sampling params are " - "accepted, which hold only when the effort resolves to 'none'." - ), - "enum": ["none", "minimal", "low", "medium", "high", "xhigh"], - }, - "comment": STRING, - "audio_transcription_config": STRING, - "vertex_ai_audio_api": { - "type": "string", - "enum": ["lyria_predict", "lyria_interactions"], - }, - } - - -def classify(key: str, modes: tuple) -> Optional[JsonSchema]: - curated = {**OBJECT_KEYS, **ARRAY_KEYS, **string_key_schemas(modes), **INTEGER_KEYS, **NUMBER_KEYS} - if key in curated: - return curated[key] - if key.startswith("supports_") or key in EXTRA_BOOLEAN_KEYS: - return BOOLEAN - if "cost" in key: - return cost_schema(key) - return None - - -def build_schema(prices: dict) -> JsonSchema: - entries = {name: entry for name, entry in prices.items() if name not in SPECIAL_ROOT_KEYS} - all_keys = tuple(sorted({key for entry in entries.values() for key in entry})) - modes = tuple(sorted({entry["mode"] for entry in entries.values() if "mode" in entry})) - unclassified = tuple(key for key in all_keys if classify(key, modes) is None) - if unclassified: - raise SystemExit( - f"Unclassified keys in {PRICES_PATH.name}: {', '.join(unclassified)}. " - f"Add them to the key tables in {Path(__file__).name} and rerun it." - ) - entry_properties = {key: classify(key, modes) for key in all_keys} - return { - "$schema": "https://json-schema.org/draft/2020-12/schema", - "title": "LiteLLM model_prices_and_context_window.json", - "description": ( - "Schema for LiteLLM's model price and context window registry " - "(https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). " - "Every top-level key except 'sample_spec' and 'fallback_generalizations' is a model id, " - "optionally prefixed with its provider (e.g. 'azure/gpt-5.4'), mapping to a model entry. " - "All costs are USD per unit. New optional fields are added regularly, so consumers should " - "ignore unknown fields rather than reject them." - ), - "type": "object", - "properties": { - "sample_spec": { - "type": "object", - "description": ( - "Documentation placeholder illustrating the entry shape; not a real model and not " - "schema-conformant (several values are prose)." - ), - }, - "fallback_generalizations": { - "type": "object", - "description": "Regex rules that generalize unknown model ids to known families; not a model entry.", - "properties": { - "rules": { - "type": "array", - "items": { - "type": "object", - "properties": { - "name": STRING, - "pattern": STRING, - "description": STRING, - }, - "required": ["name", "pattern"], - "additionalProperties": True, - }, - } - }, - "additionalProperties": False, - }, - }, - "additionalProperties": {"$ref": "#/$defs/modelEntry"}, - "$defs": { - "modelEntry": { - "type": "object", - "description": ( - "Pricing, limits, and capability flags for one model. Fields other than litellm_provider " - "are optional; boolean capability flags are simply omitted when unknown or false." - ), - "required": ["litellm_provider"], - "properties": entry_properties, - "additionalProperties": True, - } - }, - } - - -def render(schema: JsonSchema) -> str: - return json.dumps(schema, indent=2) + "\n" - - -def validation_errors(prices: dict, schema: JsonSchema) -> tuple: - validator = jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER) - return tuple( - f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}" - for error in validator.iter_errors(prices) - ) - - -def main() -> int: - check = "--check" in sys.argv[1:] - prices = json.loads(PRICES_PATH.read_text()) - rendered = render(build_schema(prices)) - errors = validation_errors(prices, json.loads(rendered)) - if errors: - print(f"{PRICES_PATH.name} does not validate against the generated schema:") - print("\n".join(errors[:20])) - return 1 - if not check: - SCHEMA_PATH.write_text(rendered) - print(f"wrote {SCHEMA_PATH}") - return 0 - if not SCHEMA_PATH.exists() or SCHEMA_PATH.read_text() != rendered: - print( - f"{SCHEMA_PATH.name} is out of sync with {PRICES_PATH.name}. " - f"Run `python {Path(__file__).relative_to(REPO_ROOT)}` and commit the result." - ) - return 1 - print(f"{SCHEMA_PATH.name} is in sync and {PRICES_PATH.name} validates against it") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/litellm-rust/crates/model-catalog/AGENTS.md b/litellm-rust/crates/model-catalog/AGENTS.md index 9fcbcd57a76..0a380788c19 100644 --- a/litellm-rust/crates/model-catalog/AGENTS.md +++ b/litellm-rust/crates/model-catalog/AGENTS.md @@ -2,5 +2,4 @@ For `model_prices_and_context_window.json` validation, we should eventually: -- Remove any schema file like `model_prices_and_context_window.schema.json` - Stop skipping this crate's tests diff --git a/tests/unit/test_cost_map_guard.py b/tests/unit/test_cost_map_guard.py index e2c54d381d8..410b3acfd42 100644 --- a/tests/unit/test_cost_map_guard.py +++ b/tests/unit/test_cost_map_guard.py @@ -22,7 +22,6 @@ def _load(name: str) -> ModuleType: return module -schema_module: Final = _load("generate_model_prices_schema") guard: Final = _load("cost_map_guard") MAP_FILES: Final = (guard.COST_MAP_PATH,) @@ -52,12 +51,9 @@ def _serialize(cost_map: dict[str, object]) -> str: return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n" -def _snapshot(cost_map: dict[str, object], backup: str | None = None, schema: str | None = None) -> object: +def _snapshot(cost_map: dict[str, object], backup: str | None = None) -> object: text = _serialize(cost_map) - rendered = schema_module.render(schema_module.build_schema(cost_map)) - return guard.Snapshot( - cost_map=text, backup=text if backup is None else backup, schema=rendered if schema is None else schema - ) + return guard.Snapshot(cost_map=text, backup=text if backup is None else backup) BASE: Final = _snapshot(BASE_MAP) @@ -135,14 +131,14 @@ def test_bot_may_add_and_reprice_models() -> None: def test_broken_json_is_reported() -> None: - head = guard.Snapshot(cost_map="{not json", backup="{not json", schema="{}") + head = guard.Snapshot(cost_map="{not json", backup="{not json") assert _failures(head, bot=False) == ( f"{guard.COST_MAP_PATH} is not valid JSON: Expecting property name enclosed in double quotes: line 1 column 2 (char 1)", ) def test_non_object_root_is_reported() -> None: - head = guard.Snapshot(cost_map="[]", backup="[]", schema="{}") + head = guard.Snapshot(cost_map="[]", backup="[]") assert _failures(head, bot=False) == (f"{guard.COST_MAP_PATH} must be a JSON object at the root",) @@ -151,28 +147,7 @@ def test_backup_drift_is_reported() -> None: assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.BACKUP_PATH)] -def test_schema_out_of_sync_is_reported() -> None: - head = _snapshot({**BASE_MAP, "openrouter/c": _entry(supports_audio_input=True)}, schema=BASE.schema) - assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.SCHEMA_PATH)] - - -def test_schema_validation_errors_are_reported() -> None: - head = _snapshot({**BASE_MAP, "openrouter/c": _entry(-1e-06)}) - prefix = f"{guard.COST_MAP_PATH} does not validate against its schema: openrouter/c." - assert [failure.removeprefix(prefix).split(":")[0] for failure in _failures(head, bot=False)] == [ - "input_cost_per_token", - "output_cost_per_token", - ] - - -def test_unclassified_entry_key_is_reported() -> None: - text = _serialize({**BASE_MAP, "openrouter/c": _entry(weird_thing=1)}) - head = guard.Snapshot(cost_map=text, backup=text, schema=BASE.schema) - (failure,) = _failures(head, bot=False) - assert "Unclassified keys" in failure and "weird_thing" in failure - - -STALE_HEAD: Final = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)}), schema="{}") +STALE_HEAD: Final = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)})) CODE_ONLY: Final = ( "litellm/utils.py", "tests/test_litellm/test_utils.py", @@ -181,7 +156,7 @@ CODE_ONLY: Final = ( def test_human_pr_that_leaves_the_cost_map_alone_skips_the_file_checks() -> None: - unparseable: Final = guard.Snapshot(cost_map="{not json", backup="", schema="") + unparseable: Final = guard.Snapshot(cost_map="{not json", backup="") assert _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=False) == () assert _failures(STALE_HEAD, changed_files=(), bot=False) == () assert _failures(unparseable, changed_files=CODE_ONLY, bot=False) == () @@ -191,18 +166,16 @@ def test_human_pr_that_leaves_the_cost_map_alone_skips_the_file_checks() -> None def test_touching_any_cost_map_file_keeps_the_file_checks(guarded_path: str) -> None: failures: Final = _failures(STALE_HEAD, changed_files=(*CODE_ONLY, guarded_path), bot=False) assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)] - assert [failure for failure in failures if failure.startswith(guard.SCHEMA_PATH)] def test_bot_pr_always_gets_the_file_checks() -> None: failures: Final = _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=True) assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)] - assert [failure for failure in failures if failure.startswith(guard.SCHEMA_PATH)] def test_contract_names_the_skip() -> None: assert guard.contract_for(False, CODE_ONLY) == "human PR, cost map untouched" - assert guard.contract_for(False, (*CODE_ONLY, guard.SCHEMA_PATH)) == "human PR, file checks only" + assert guard.contract_for(False, (*CODE_ONLY, guard.COST_MAP_PATH)) == "human PR, file checks only" assert guard.contract_for(True, CODE_ONLY) == "bot contract enforced" @@ -215,6 +188,13 @@ def test_bot_may_only_touch_the_cost_map_files() -> None: ) +def test_bot_may_update_cost_map_and_legacy_schema() -> None: + head: Final = _snapshot({**BASE_MAP, "openrouter/c": _entry()}) + changed: Final = (*guard.GUARDED_PATHS, guard.LEGACY_SCHEMA_PATH) + + assert _failures(head, changed_files=changed) == () + + def test_bot_may_not_remove_models() -> None: head = _snapshot({key: value for key, value in BASE_MAP.items() if key != "openrouter/b"}) assert _failures(head, bot=False) == () @@ -238,7 +218,6 @@ def _commit(repo: Path, cost_map: dict[str, object], message: str) -> str: (repo / guard.COST_MAP_PATH).write_text(text) (repo / guard.BACKUP_PATH).parent.mkdir(exist_ok=True) (repo / guard.BACKUP_PATH).write_text(text) - (repo / guard.SCHEMA_PATH).write_text(schema_module.render(schema_module.build_schema(cost_map))) return _git_commit(repo, message) @@ -292,8 +271,7 @@ def test_main_skips_the_file_checks_on_a_stale_base_the_pr_never_touched(tmp_pat subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) _commit(tmp_path, BASE_MAP, "base") (tmp_path / guard.BACKUP_PATH).write_text(_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)})) - (tmp_path / guard.SCHEMA_PATH).write_text("{}") - stale_base: Final = _commit_code_only(tmp_path, "stale base with drifted backup and schema") + stale_base: Final = _commit_code_only(tmp_path, "stale base with drifted backup") head: Final = _commit_code_only(tmp_path, "code change on the stale base") human: Final = _run_guard(tmp_path, stale_base, head, "litellm_fix_pricing") assert human.returncode == 0, human.stdout + human.stderr diff --git a/tests/unit/test_model_prices_schema.py b/tests/unit/test_model_prices_and_context_window.py similarity index 67% rename from tests/unit/test_model_prices_schema.py rename to tests/unit/test_model_prices_and_context_window.py index 211a43d2574..49a6fb05b15 100644 --- a/tests/unit/test_model_prices_schema.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -1,6 +1,5 @@ from __future__ import annotations -import importlib.util import json import re from collections.abc import Mapping @@ -8,7 +7,6 @@ from pathlib import Path from types import MappingProxyType from typing import Final -import jsonschema import pytest import litellm @@ -17,27 +15,8 @@ from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts REPO_ROOT = Path(__file__).parents[2] -GENERATOR_PATH = REPO_ROOT / "ci_cd" / "generate_model_prices_schema.py" PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json" BACKUP_PRICES_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" -SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json" - - -def build_validator(schema: dict) -> jsonschema.Draft202012Validator: - return jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER) - - -def load_generator(): - spec = importlib.util.spec_from_file_location("generate_model_prices_schema", GENERATOR_PATH) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -@pytest.fixture(scope="module") -def committed_schema() -> dict: - return json.loads(SCHEMA_PATH.read_text()) @pytest.fixture(scope="module") @@ -45,183 +24,6 @@ def prices() -> dict: return json.loads(PRICES_PATH.read_text()) -def test_committed_schema_matches_generator_output(prices: dict, committed_schema: dict): - generator = load_generator() - regenerated = json.loads(generator.render(generator.build_schema(prices))) - assert regenerated == committed_schema, ( - "model_prices_and_context_window.schema.json is out of sync; " - "run `python ci_cd/generate_model_prices_schema.py` and commit the result" - ) - - -def test_prices_file_validates_against_committed_schema(prices: dict, committed_schema: dict): - validator = build_validator(committed_schema) - errors = [ - f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}" - for error in validator.iter_errors(prices) - ] - assert errors == [] - - -@pytest.mark.parametrize( - "entry", - [ - {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": "0.01"}, - {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": -1}, - {"litellm_provider": "openai", "mode": "not_a_real_mode"}, - {"mode": "chat"}, - {"litellm_provider": "openai", "deprecation_date": "June 2026"}, - {"litellm_provider": "openai", "deprecation_date": "2026-99-99"}, - {"litellm_provider": "openai", "deprecation_date": "2026-13-01"}, - {"litellm_provider": "openai", "deprecation_date": "2026-01-32"}, - {"litellm_provider": "openai", "deprecation_date": "2026-01-00"}, - {"litellm_provider": "openai", "deprecation_date": "2026-02-31"}, - {"litellm_provider": "openai", "supported_modalities": ["smell"]}, - {"litellm_provider": "openai", "supports_vision": "yes"}, - {"litellm_provider": "openai", "max_tokens": 8191.5}, - {"litellm_provider": "openai", "tiered_pricing": [{"unknown_tier_field": 1}]}, - ], - ids=[ - "cost_as_string", - "negative_cost", - "unknown_mode", - "missing_provider", - "non_iso_deprecation_date", - "impossible_month_and_day", - "month_out_of_range", - "day_out_of_range", - "day_zero", - "calendar_impossible_day", - "unknown_modality", - "boolean_flag_as_string", - "fractional_max_tokens", - "unknown_tiered_pricing_field", - ], -) -def test_schema_rejects_malformed_entries(committed_schema: dict, entry: dict): - validator = build_validator(committed_schema) - assert not validator.is_valid({"some-model": entry}) - - -def test_schema_accepts_minimal_and_unknown_optional_fields(committed_schema: dict): - validator = build_validator(committed_schema) - assert validator.is_valid({"some-model": {"litellm_provider": "openai"}}) - assert validator.is_valid({"some-model": {"litellm_provider": "openai", "brand_new_field": {"nested": True}}}) - - -def test_schema_accepts_cache_creation_cost_inside_a_pricing_tier(committed_schema: dict): - validator = build_validator(committed_schema) - entry = { - "litellm_provider": "dashscope", - "mode": "chat", - "tiered_pricing": [ - { - "range": [0, 256000], - "input_cost_per_token": 3.25e-07, - "output_cost_per_token": 1.95e-06, - "cache_creation_input_token_cost": 4.063e-07, - "cache_read_input_token_cost": 3.25e-08, - } - ], - } - assert validator.is_valid({"some-model": entry}) - - -OFF_PEAK_ENTRY: Final = MappingProxyType( - { - "litellm_provider": "openrouter", - "mode": "chat", - "input_cost_per_token": 2e-6, - "output_cost_per_token": 8e-6, - "off_peak_pricing": { - "hours_utc": "16:30-00:30", - "windows": [{"hours_utc": ["00:30-02:00"], "weekdays": [6, "Sunday", "mon", "THURS"]}], - "weekday_timezone": "Asia/Shanghai", - "input_cost_per_token": 1e-6, - "output_cost_per_token": 4e-6, - "cache_read_input_token_cost": 1e-7, - }, - } -) - - -def test_generator_classifies_off_peak_pricing_as_a_windowed_rate_block(): - generator = load_generator() - schema = json.loads(generator.render(generator.build_schema({"some-model": dict(OFF_PEAK_ENTRY)}))) - validator = build_validator(schema) - assert validator.is_valid({"some-model": dict(OFF_PEAK_ENTRY)}) - - -@pytest.mark.parametrize( - "block", - [ - {"hours_utc": "16:30-00:30", "input_cost_per_token": "1e-6"}, - {"hours_utc": "16:30-00:30", "input_cost_per_token": -1e-6}, - {"hours_utc": 1630, "input_cost_per_token": 1e-6}, - {"hours_utc": "16:30-00:30", "discount": 0.5}, - {"windows": [{"weekdays": [6]}], "input_cost_per_token": 1e-6}, - {"windows": [{"hours_utc": "00:30-02:00", "weekdays": [0]}], "input_cost_per_token": 1e-6}, - {"windows": [], "input_cost_per_token": 1e-6}, - {"input_cost_per_token": 1e-6}, - {"hours_utc": "16:30", "input_cost_per_token": 1e-6}, - {"hours_utc": "25:00-01:00", "input_cost_per_token": 1e-6}, - {"hours_utc": ["16:30-00:30", "4pm-midnight"], "input_cost_per_token": 1e-6}, - {"windows": [{"hours_utc": "00:30-02:00", "weekdays": ["Funday"]}], "input_cost_per_token": 1e-6}, - ], -) -def test_generated_off_peak_schema_rejects_malformed_blocks(block: dict): - generator = load_generator() - schema = json.loads(generator.render(generator.build_schema({"some-model": dict(OFF_PEAK_ENTRY)}))) - validator = build_validator(schema) - assert not validator.is_valid({"some-model": {**OFF_PEAK_ENTRY, "off_peak_pricing": block}}) - - -TOKEN_COUNT_ENTRY: Final = MappingProxyType( - { - "litellm_provider": "bedrock_converse", - "mode": "chat", - "token_count_framing": 78, - "token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"}, - } -) - - -@pytest.mark.parametrize( - "setting", - [ - {"token_count_framing": 0}, - {"token_count_framing": 78}, - {"token_count_model": {"model": "us.anthropic.claude-opus-5-5"}}, - ], -) -def test_generator_classifies_token_count_settings(setting: dict): - generator = load_generator() - schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) - assert build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}}) - - -@pytest.mark.parametrize( - "setting", - [ - {"token_count_framing": -1}, - {"token_count_framing": "78"}, - {"token_count_framing": 1.5}, - {"token_count_model": {"model": ""}}, - {"token_count_model": {"model": " "}}, - {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": ""}}, - {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": " "}}, - {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": 5}}, - {"token_count_model": {"region": "us-east-1"}}, - {"token_count_model": "us.anthropic.claude-opus-5-5"}, - {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "zone": "us-east-1"}}, - ], -) -def test_generated_token_count_schema_rejects_malformed_settings(setting: dict): - generator = load_generator() - schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) - assert not build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}}) - - def find_duplicate_keys(path: Path) -> list[str]: duplicates: list[str] = [] From 748e03088d0415076d4342061002c227113afbec Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:54:57 -0700 Subject: [PATCH 21/36] fix(count_tokens)!: answer 400 for Vertex partner models without a count API and OpenAI-shaped provider errors on /v1/responses/input_tokens (#47163) Backport of #47163 to stable/1.104.x. Cherry-picked from ef1d3959a06289eae04b89c94426e7c93ebbb9a5 (main). --- litellm/llms/vertex_ai/common_utils.py | 18 ++ .../count_tokens/handler.py | 36 ++- .../common_utils/openai_error_payload.py | 4 + .../proxy/response_api_endpoints/endpoints.py | 97 +++++++- .../test_openai_count_tokens_tools_wire.py | 25 ++ .../test_vertex_partner_count_tokens_wire.py | 177 ++++++++++++++- .../response_api_endpoints/test_endpoints.py | 214 ++++++++++++++++++ .../vertex_ai/test_vertex_ai_common_utils.py | 29 +++ tests/unit/proxy/test_custom_tokenizer_bug.py | 60 ++++- 9 files changed, 628 insertions(+), 32 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 2c7e07616b5..eb57ef4e6c6 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1242,6 +1242,7 @@ class VertexAITokenCounter(BaseTokenCounter): ) -> TokenCountResponse | None: import copy + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( VertexAIError as PartnerVertexAIError, ) @@ -1252,6 +1253,23 @@ class VertexAITokenCounter(BaseTokenCounter): deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) + without_count_route: Final = ( + VertexAIPartnerModels.is_vertex_partner_model(model_to_use) and count_tokens_publisher(model_to_use) is None + ) + model_info: Final[Mapping[str, object]] = deployment.get("model_info") or {} + if without_count_route and messages and model_info.get("custom_tokenizer") is not None: + return None + if without_count_route: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="vertex_ai_partner_models", + error=True, + error_message=f"LiteLLM does not support counting tokens for Vertex AI model {model_to_use}.", + status_code=400, + ) + # Check if this is a partner model (Claude, Mistral, etc.) if VertexAIPartnerModels.is_vertex_partner_model(model_to_use): # Use partner models token counter diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 285380f5fc7..7100c9efde3 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -1,5 +1,5 @@ """ -Token counter for Vertex AI Partner Models (Anthropic Claude, Mistral, etc.) +Token counter for Vertex AI Partner Models (Anthropic Claude, Meta Llama) This handler provides token counting for partner models hosted on Vertex AI. Unlike Gemini models which use Google's token counting API, partner models use @@ -12,36 +12,30 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +_COUNT_TOKENS_PUBLISHERS: Final = ( + ("claude", "anthropic"), + ("llama", "meta"), + ("meta/", "meta"), +) + + +def count_tokens_publisher(model: str) -> str | None: + return next((publisher for marker, publisher in _COUNT_TOKENS_PUBLISHERS if marker in model), None) + class VertexAIPartnerModelsTokenCounter(VertexBase): """ Token counter for Vertex AI Partner Models. - Handles token counting for models like Claude (Anthropic), Mistral, etc. + Handles token counting for models like Claude (Anthropic) and Llama (Meta) that are available through Vertex AI's partner model program. """ def _get_publisher_for_model(self, model: str) -> str: - """ - Determine the publisher name for the given model. - - Args: - model: The model name (e.g., "claude-3-5-sonnet-20241022") - - Returns: - Publisher name to use in the Vertex AI endpoint URL - - Raises: - ValueError: If the model is not a recognized partner model - """ - if "claude" in model: - return "anthropic" - elif "mistral" in model or "codestral" in model: - return "mistralai" - elif "llama" in model or "meta/" in model: - return "meta" - else: + publisher: Final = count_tokens_publisher(model) + if publisher is None: raise ValueError(f"Unknown partner model: {model}") + return publisher def _build_count_tokens_endpoint( self, diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index 202c61b620e..b24761ac039 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -42,6 +42,10 @@ def openai_error_type(exc: object, status_code: int) -> str: carried: Final = attribute_of(exc, "type") if isinstance(carried, str) and carried != STRINGIFIED_NONE: return carried + return openai_error_type_for_status(status_code) + + +def openai_error_type_for_status(status_code: int) -> str: mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) if mapped is not None: return mapped diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 75eefb2e73b..021c741b2bc 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -14,11 +14,12 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse from openai.types.responses import ResponseItemList from openai.types.responses.response_create_params import ResponseInputParam -from pydantic import BaseModel, ConfigDict, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, with_config from starlette.websockets import WebSocket, WebSocketDisconnect -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping from litellm.constants import EMPTY_MAPPING from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -31,10 +32,16 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth_websocket, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_set_request_parsed_body, ) +from litellm.proxy.common_utils.openai_error_payload import ( + LITELLM_CALL_ID_HEADER, + error_status_code, + openai_error_type_for_status, +) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -1103,6 +1110,85 @@ def _responses_invalid_request_response(message: str, param: str | None, code: s return JSONResponse(status_code=400, content=body) +class _ProviderErrorMessage(TypedDict): + message: ReadOnly[str] + + +class _NestedProviderError(TypedDict): + error: ReadOnly[_ProviderErrorMessage] + + +@with_config(ConfigDict(extra="ignore")) +class _ProviderOpenAIErrorDetail(TypedDict): + message: ReadOnly[str] + type: ReadOnly[str] + param: NotRequired[ReadOnly[str | None]] + code: NotRequired[ReadOnly[str | None]] + + +@with_config(ConfigDict(extra="forbid")) +class _ProviderOpenAIError(TypedDict): + error: ReadOnly[_ProviderOpenAIErrorDetail] + + +_PROVIDER_COUNT_ERROR_TYPE: Final = "token_counting_error" +_PROVIDER_OPENAI_ERROR: Final = TypeAdapter(_ProviderOpenAIError) +_NESTED_PROVIDER_ERROR: Final = TypeAdapter(_NestedProviderError) + + +def _provider_error_json(raw_message: str) -> str: + start: Final = raw_message.find("{") + return raw_message if start == -1 else raw_message[start:] + + +def _provider_error_message(raw_message: str) -> str: + error_json: Final = _provider_error_json(raw_message) + try: + return _NESTED_PROVIDER_ERROR.validate_json(error_json)["error"]["message"] + except ValidationError: + flat_message: Final = AnthropicExceptionMapping.extract_error_message(error_json) + return raw_message if flat_message == error_json else flat_message + + +def _provider_count_error_body(raw_message: str, status_code: int) -> _ResponsesApiErrorBody: + try: + provider_error: Final = _PROVIDER_OPENAI_ERROR.validate_json(_provider_error_json(raw_message))["error"] + except ValidationError: + return { + "error": { + "message": _provider_error_message(raw_message), + "type": openai_error_type_for_status(status_code), + "param": None, + "code": None, + } + } + return { + "error": { + "message": provider_error["message"], + "type": provider_error["type"], + "param": provider_error.get("param"), + "code": provider_error.get("code"), + } + } + + +def _provider_count_error_response(exc: ProxyException, request: Request) -> JSONResponse: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; a route that answers an error itself must stamp the OTel server span like the global ProxyException handler does + general_settings_view, + ) + + status_code: Final = error_status_code(exc, 500) + _close_dangling_otel_server_span(request, status_code, exc=exc) + body: Final = _provider_count_error_body(exc.message, status_code) + call_id: Final = error_body_call_id(general_settings_view(), exc.headers.get(LITELLM_CALL_ID_HEADER)) + return JSONResponse( + status_code=status_code, + content={"error": with_call_id({**body["error"]}, call_id)}, + headers=exc.headers, + ) + + def _missing_responses_param_response(param: str) -> JSONResponse: return _responses_invalid_request_response( message=f"Missing required parameter: '{param}'.", @@ -1194,7 +1280,12 @@ async def responses_input_tokens( message=f"Invalid request for token counting: {e}", param=None, code=None ) - token_response: Final = await token_counter(request=token_request, call_endpoint=True) + try: + token_response: Final = await token_counter(request=token_request, call_endpoint=True) + except ProxyException as e: + if e.type != _PROVIDER_COUNT_ERROR_TYPE: + raise + return _provider_count_error_response(e, request) result: Final[_ResponsesInputTokensResult] = { "object": "response.input_tokens", "input_tokens": token_response.total_tokens, diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index 2b23ccdf4f2..16a5b8685ad 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -391,3 +391,28 @@ def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: O ) assert response.status_code == 200, response.text assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English." + + +@pytest.mark.parametrize( + "error", + [ + pytest.param( + {"message": "Invalid input", "type": "invalid_request_error", "param": "input", "code": "invalid_value"}, + id="full", + ), + pytest.param({"message": "Invalid input", "type": "invalid_request_error"}, id="bare"), + ], +) +def test_responses_input_tokens_passes_an_openai_count_error_through( + counting_proxy: OwnedProxy, error: dict[str, JsonValue] +) -> None: + def rejecting(_request: Request) -> Reply: + return Reply(status=400, body=json.dumps({"error": error}).encode()) + + gateway: Final = counting_proxy.gateway + with wire_server(rejecting) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/responses/input_tokens", {"model": model, "input": "hello"}) + assert response.status_code == 400, response.text + assert response.json() == {"error": {"param": None, "code": None, **error}}, response.text + assert _count_body(wire)["input"] == [{"role": "user", "content": "hello"}] diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index ae9b8434d11..1c4b1d1a6e5 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -7,8 +7,10 @@ import uuid from collections.abc import Callable, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack +from dataclasses import dataclass from pathlib import Path from queue import SimpleQueue +from types import MappingProxyType from typing import Final import httpx @@ -16,10 +18,11 @@ import psutil import pytest import yaml from integration._support.client import Gateway, Scenario, eventually -from integration._support.process import owned_proxy_process +from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +from tokenizers import Regex, Tokenizer, models, pre_tokenizers _BACKEND: Final = "claude-sonnet-4-6" _PROJECT: Final = "scripted-project" @@ -35,9 +38,19 @@ _REPLY_TEXT: Final = "scripted reply" _OWNED_MODEL: Final = "partner-claude" _OWNED_UNREACHABLE_MODEL: Final = "partner-claude-unreachable" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_GLM_BACKEND: Final = "zai-org/glm-5.2-maas" +_MISTRAL_BACKEND: Final = "mistral-medium-3" +_LLAMA_BACKEND: Final = "meta/llama-4-maverick-17b-128e-instruct-maas" +_LLAMA_COUNT_TARGET: Final = ( + f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/meta/models/count-tokens:rawPredict" +) +_COUNTED_TEXT: Final = "Count this message" +_TOKENIZER_REPO: Final = "scripted/char-tokenizer" +_TOKENIZER_TARGET: Final = f"/{_TOKENIZER_REPO}/resolve/main/tokenizer.json" +_TOKENIZER_HEADERS: Final = MappingProxyType({"X-Repo-Commit": "0" * 40, "ETag": '"scripted-char-tokenizer"'}) _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") -_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this message"}] +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": _COUNTED_TEXT}] _SYSTEM: Final = "You are a terse assistant that answers in one sentence" _SYSTEM_BLOCKS: Final[list[JsonValue]] = [ {"type": "text", "text": "You are a terse assistant"}, @@ -115,16 +128,16 @@ def _counted(_request: Request) -> Reply: return Reply(body=json.dumps({"input_tokens": _PEER_COUNT}).encode()) -def _rejected(status: int) -> Reply: +def _rejected(status: int, error_type: str = "invalid_request_error") -> Reply: return Reply( status=status, - body=json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": _REJECTION}}).encode(), + body=json.dumps({"type": "error", "error": {"type": error_type, "message": _REJECTION}}).encode(), ) -def _rejecting(status: int) -> Callable[[Request], Reply]: +def _rejecting(status: int, error_type: str = "invalid_request_error") -> Callable[[Request], Reply]: def count(_request: Request) -> Reply: - return _rejected(status) + return _rejected(status, error_type) return count @@ -224,6 +237,78 @@ def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int: return total +@dataclass(frozen=True, slots=True) +class _CountRoute: + path: Callable[[str], str] + body: Callable[[str], dict[str, JsonValue]] + rejection: Callable[[str], dict[str, JsonValue]] + count_field: str + params: Mapping[str, str] = MappingProxyType({}) + + def send(self, gateway: Gateway, model: str) -> httpx.Response: + return gateway.request("POST", self.path(model), self.body(model), params=self.params) + + +_MESSAGES_ROUTE: Final = _CountRoute( + path=lambda _model: "/v1/messages/count_tokens", + body=_bare, + rejection=lambda message: {"type": "error", "error": {"type": "invalid_request_error", "message": message}}, + count_field="input_tokens", +) +_RESPONSES_ROUTE: Final = _CountRoute( + path=lambda _model: "/v1/responses/input_tokens", + body=lambda model: {"model": model, "input": _COUNTED_TEXT}, + rejection=lambda message: { + "error": {"message": message, "type": "invalid_request_error", "param": None, "code": None} + }, + count_field="input_tokens", +) +_GOOGLE_ROUTE: Final = _CountRoute( + path=lambda model: f"/v1beta/models/{model}:countTokens", + body=lambda _model: {"contents": [{"role": "user", "parts": [{"text": _COUNTED_TEXT}]}]}, + rejection=lambda message: {"error": {"code": 400, "message": message, "status": "INVALID_ARGUMENT"}}, + count_field="totalTokens", +) +_UTILS_ROUTE: Final = _CountRoute( + path=lambda _model: "/utils/token_counter", + body=_bare, + rejection=lambda message: { + "error": {"message": message, "type": "token_counting_error", "param": "model", "code": "400"} + }, + count_field="total_tokens", + params=MappingProxyType({"call_endpoint": "true"}), +) +_COUNT_ROUTES: Final = (_MESSAGES_ROUTE, _RESPONSES_ROUTE, _GOOGLE_ROUTE, _UTILS_ROUTE) + + +def _error_message(payload: Mapping[str, JsonValue]) -> str: + error: Final = payload.get("error") + assert isinstance(error, dict), payload + message: Final = error.get("message") + assert isinstance(message, str), payload + return message + + +def _route_count(gateway: Gateway, route: _CountRoute, model: str) -> tuple[int, JsonValue]: + response: Final = route.send(gateway, model) + return response.status_code, _JSON_OBJECT.validate_json(response.content).get(route.count_field) + + +def _char_tokenizer() -> bytes: + tokenizer: Final = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) + tokenizer.pre_tokenizer = pre_tokenizers.Split(Regex("."), behavior="isolated") + return tokenizer.to_str().encode() + + +def _tokenizer_hub(tokenizer: bytes) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == _TOKENIZER_TARGET: + return Reply(body=tokenizer, headers=_TOKENIZER_HEADERS) + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + return respond + + def _closed_port_url() -> str: with socket.socket() as reserve: reserve.bind(("127.0.0.1", 0)) @@ -374,11 +459,89 @@ def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> Non {"model": model, "input": "Count this message", "instructions": "Be terse", "tools": _RESPONSES_TOOLS}, ) assert response.status_code == 400, response.text - assert "input_tokens" not in response.text, response.text + assert response.json() == _RESPONSES_ROUTE.rejection(_REJECTION), response.text (sent,) = _counted_bodies(wire) assert sent["tools"] == _RESPONSES_TOOLS, sent +def test_responses_input_tokens_types_a_peer_error_by_its_status(gateway: Gateway) -> None: + with wire_server(_peer(_rejecting(404, "not_found_error"))) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _RESPONSES_ROUTE.send(gateway, model) + assert response.status_code == 404, response.text + assert response.json() == _RESPONSES_ROUTE.rejection(_REJECTION), response.text + assert _counted_bodies(wire) == (_PEER_BARE,) + + +@pytest.mark.parametrize( + ("backend", "route"), + [ + pytest.param(_GLM_BACKEND, _MESSAGES_ROUTE, id="glm-messages"), + pytest.param(_GLM_BACKEND, _RESPONSES_ROUTE, id="glm-responses"), + pytest.param(_GLM_BACKEND, _GOOGLE_ROUTE, id="glm-google"), + pytest.param(_GLM_BACKEND, _UTILS_ROUTE, id="glm-utils"), + pytest.param(_MISTRAL_BACKEND, _MESSAGES_ROUTE, id="mistral-messages"), + ], +) +def test_partner_without_a_count_route_answers_400_without_calling_vertex( + gateway: Gateway, backend: str, route: _CountRoute +) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{backend}") + response: Final = route.send(gateway, model) + assert response.status_code == 400, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + message: Final = _error_message(payload) + assert backend in message, message + assert payload == route.rejection(message), response.text + assert wire.drain() == () + + +def test_llama_count_still_reaches_the_meta_count_route(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{_LLAMA_BACKEND}") + response: Final = _count(gateway, _bare(model)) + assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + (sent,) = _count_bodies(wire.drain(), _LLAMA_COUNT_TARGET) + assert sent["messages"] == _MESSAGES, sent + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_custom_tokenizer_partner_counts_locally_on_every_count_route(gateway: Gateway, tmp_path: Path) -> None: + environment: Final = {"HF_HOME": str(tmp_path / "hf-home")} + with ( + wire_server(_tokenizer_hub(_char_tokenizer())) as hub, + wire_server(_peer()) as vertex, + owned_proxy_process(gateway, tmp_path, {**environment, "HF_ENDPOINT": hub.url}) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = _deployment( + owned.gateway, + scenario, + vertex.url, + model=f"vertex_ai/{_GLM_BACKEND}", + model_info={"custom_tokenizer": {"identifier": _TOKENIZER_REPO, "revision": "main", "auth_token": None}}, + ) + local: Final = _payload( + owned.gateway.request("POST", "/utils/token_counter", _bare(model), params={"call_endpoint": "false"}) + ) + assert local["tokenizer_type"] == "huggingface_tokenizer", local + total: Final = local["total_tokens"] + assert isinstance(total, int) and total >= len(_COUNTED_TEXT), local + counts: Final = tuple(_route_count(owned.gateway, route, model) for route in _COUNT_ROUTES) + assert counts == ((200, total),) * len(_COUNT_ROUTES), counts + contents_only: Final = owned.gateway.request( + "POST", + "/utils/token_counter", + {"model": model, "contents": [{"role": "user", "parts": [{"text": _COUNTED_TEXT}]}]}, + params={"call_endpoint": "true"}, + ) + assert contents_only.status_code == 400, contents_only.text + assert _GLM_BACKEND in _error_message(_JSON_OBJECT.validate_json(contents_only.content)), contents_only.text + assert vertex.drain() == () + assert {request.target for request in hub.drain()} >= {_TOKENIZER_TARGET} + + def test_gemini_count_tokens_route_reaches_the_partner_peer(gateway: Gateway) -> None: with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index e684aa55b33..d7648a08b6a 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -3,6 +3,7 @@ Test for response_api_endpoints/endpoints.py """ import unittest +from collections.abc import Mapping from typing import Any, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -2428,6 +2429,219 @@ class TestResponsesInputTokens: assert response.status_code == 429, response.text assert response.json()["error"]["message"] == "rate limited" + @pytest.mark.parametrize( + "provider_body, status_code, expected_error", + [ + pytest.param( + { + "error": { + "message": "Missing required parameter: 'tools[0].name'.", + "type": "invalid_request_error", + "param": "tools[0].name", + "code": "missing_required_parameter", + } + }, + 400, + { + "message": "Missing required parameter: 'tools[0].name'.", + "type": "invalid_request_error", + "param": "tools[0].name", + "code": "missing_required_parameter", + }, + id="openai_error_passes_through", + ), + pytest.param( + { + "error": { + "message": "Missing required parameter: 'tools[0].name'.", + "type": "invalid_request_error", + "param": "tools[0].name", + } + }, + 400, + { + "message": "Missing required parameter: 'tools[0].name'.", + "type": "invalid_request_error", + "param": "tools[0].name", + "code": None, + }, + id="openai_error_without_code_keeps_type_and_param", + ), + pytest.param( + { + "type": "error", + "error": {"type": "invalid_request_error", "message": "tools.0: Input tag 'function' is invalid"}, + "request_id": "req_vrtx_1", + }, + 400, + { + "message": "tools.0: Input tag 'function' is invalid", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="anthropic_error_message_unwrapped", + ), + pytest.param( + { + "type": "error", + "error": {"type": "not_found_error", "message": "model: claude-typo"}, + "request_id": "req_vrtx_2", + }, + 404, + {"message": "model: claude-typo", "type": "invalid_request_error", "param": None, "code": None}, + id="anthropic_error_type_follows_the_status", + ), + pytest.param( + {"error": {"code": 404, "message": "Publisher model not found.", "status": "NOT_FOUND"}}, + 404, + {"message": "Publisher model not found.", "type": "invalid_request_error", "param": None, "code": None}, + id="google_error_message_unwrapped", + ), + pytest.param( + "litellm.APIError: Google Gen AI Studio API error: 404 - " + + json.dumps( + {"error": {"code": 404, "message": "models/gemini-typo is not found.", "status": "NOT_FOUND"}} + ), + 404, + { + "message": "models/gemini-typo is not found.", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="prefixed_gemini_error_message_unwrapped", + ), + pytest.param( + "litellm.BadRequestError: AzureException - " + + json.dumps( + { + "error": { + "message": "Unrecognized request argument supplied: foo", + "type": "invalid_request_error", + "param": "foo", + "code": "unknown_parameter", + } + } + ), + 400, + { + "message": "Unrecognized request argument supplied: foo", + "type": "invalid_request_error", + "param": "foo", + "code": "unknown_parameter", + }, + id="prefixed_openai_error_passes_through", + ), + pytest.param( + {"message": "The provided model doesn't support counting tokens."}, + 400, + { + "message": "The provided model doesn't support counting tokens.", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="bedrock_error_message_unwrapped", + ), + pytest.param( + "BedrockException - " + json.dumps({"message": "The provided model doesn't support counting tokens."}), + 400, + { + "message": "The provided model doesn't support counting tokens.", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="prefixed_bedrock_error_message_unwrapped", + ), + pytest.param( + "tool schema {unclosed is invalid", + 400, + { + "message": "tool schema {unclosed is invalid", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="plain_text_with_a_brace_kept_whole", + ), + pytest.param( + "upstream count service unavailable", + 503, + { + "message": "upstream count service unavailable", + "type": "internal_server_error", + "param": None, + "code": None, + }, + id="plain_text_error", + ), + ], + ) + def test_provider_count_error_answers_in_openai_error_shape( + self, + provider_body: str | Mapping[str, object], + status_code: int, + expected_error: Mapping[str, str | None], + ) -> None: + from litellm.proxy._types import ProxyException + + failing_counter: Final = AsyncMock( + side_effect=ProxyException( + message=provider_body if isinstance(provider_body, str) else json.dumps(provider_body), + type="token_counting_error", + param="model", + code=str(status_code), + ) + ) + response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) + + assert response.status_code == status_code, response.text + assert response.json() == {"error": expected_error} + + def test_non_provider_count_error_keeps_its_own_type(self) -> None: + from litellm.proxy._types import ProxyException + + failing_counter: Final = AsyncMock( + side_effect=ProxyException( + message="Token counting is disabled and no provider API result available", + type="token_counting_disabled", + param="model", + code=503, + ) + ) + response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) + + assert response.status_code == 503, response.text + assert response.json()["error"]["type"] == "token_counting_disabled" + + def test_provider_count_error_names_the_call_id_when_opted_in(self) -> None: + from litellm.proxy._types import ProxyException + + failing_counter: Final = AsyncMock( + side_effect=ProxyException( + message="LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas.", + type="token_counting_error", + param="model", + code="400", + headers={"x-litellm-call-id": "call-9564"}, + ) + ) + with patch("litellm.proxy.proxy_server.general_settings", {"include_call_id_in_error_body": True}): + response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) + + assert response.status_code == 400, response.text + assert response.json() == { + "error": { + "message": "LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas.", + "type": "invalid_request_error", + "param": None, + "code": None, + "litellm_call_id": "call-9564", + } + } + def test_responses_routes_document_response_models_in_openapi_schema(): from typing import cast diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index f9561a0d710..71e6ecc1f81 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -47,6 +47,35 @@ async def test_vertex_gemini_rejects_non_text_system_blocks_before_request( assert "text blocks" in (result.error_message or "") +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", + ( + "zai-org/glm-5.2-maas", + "deepseek-ai/deepseek-v3.2-maas", + "qwen/qwen3-coder-480b-a35b-instruct-maas", + "mistral-medium-3", + "codestral-2501", + "openai/gpt-oss-120b-maas", + "google/gemma-4-26b-a4b-it-maas", + ), +) +async def test_vertex_partner_model_without_an_inference_matching_count_answers_400_naming_the_model(model: str) -> None: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use=model, + messages=[{"role": "user", "content": "Count this"}], + contents=None, + deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "global"}}, + request_model=f"vertex_ai/{model}", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert result.error_message == f"LiteLLM does not support counting tokens for Vertex AI model {model}." + + class _VertexCountTokensPeer: def __init__(self, status_code: int) -> None: self.status_code: Final = status_code diff --git a/tests/unit/proxy/test_custom_tokenizer_bug.py b/tests/unit/proxy/test_custom_tokenizer_bug.py index 23b6b302202..29625ea1a32 100644 --- a/tests/unit/proxy/test_custom_tokenizer_bug.py +++ b/tests/unit/proxy/test_custom_tokenizer_bug.py @@ -9,6 +9,7 @@ The HuggingFace download boundary (Tokenizer.from_pretrained) is mocked so these stay hermetic unit tests; the proxy's extraction-and-selection path runs for real. """ +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -17,7 +18,7 @@ import litellm import litellm.proxy.proxy_server import litellm.utils from litellm import Router -from litellm.proxy._types import TokenCountRequest +from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.proxy_server import token_counter @@ -102,3 +103,60 @@ async def test_model_without_custom_tokenizer_uses_default(monkeypatch): assert response.tokenizer_type == "openai_tokenizer" assert response.model_used == "gpt-4" assert response.total_tokens > 0 + + + +def _vertex_glm_router_with_its_own_tokenizer() -> Router: + return Router( + model_list=[ + { + "model_name": "vertex-glm", + "litellm_params": { + "model": "vertex_ai/zai-org/glm-5.2-maas", + "vertex_project": "test-project", + "vertex_location": "global", + }, + "model_info": { + "custom_tokenizer": {"identifier": "zai-org/GLM-5.2", "revision": "main", "auth_token": None}, + }, + } + ] + ) + + +@pytest.mark.asyncio +async def test_custom_tokenizer_counts_a_vertex_partner_model_without_a_count_route( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _vertex_glm_router_with_its_own_tokenizer()) + + with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls: + mock_tokenizer_cls.from_pretrained.return_value = _fake_hf_tokenizer(9) + + response: Final = await token_counter( + request=TokenCountRequest(model="vertex-glm", messages=[{"role": "user", "content": "Return exactly OK."}]), + call_endpoint=True, + ) + + mock_tokenizer_cls.from_pretrained.assert_called_once_with("zai-org/GLM-5.2", revision="main", token=None) + assert response.error is not True + assert response.tokenizer_type == "huggingface_tokenizer" + assert response.total_tokens >= 9 + + +@pytest.mark.asyncio +async def test_contents_only_count_of_a_vertex_partner_model_without_a_count_route_answers_400( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _vertex_glm_router_with_its_own_tokenizer()) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest( + model="vertex-glm", contents=[{"role": "user", "parts": [{"text": "Return exactly OK."}]}] + ), + call_endpoint=True, + ) + + assert raised.value.code == "400" + assert raised.value.message == "LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas." From 52de3b0a29d196f281687f7c53c963aa429e7e86 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:55:36 -0700 Subject: [PATCH 22/36] test(integration): add the support helpers the backported count_tokens wire tests import Hand-ported to stable/1.104.x from main (bedrock_runtime_peer, provider, _cache_control_marks_support verbatim, graceful_stop_seconds at the line's 30s stop wait). --- .../_support/bedrock_runtime_peer.py | 278 ++++++++++ tests/integration/_support/process.py | 6 +- tests/integration/_support/provider.py | 49 ++ .../providers/_cache_control_marks_support.py | 479 ++++++++++++++++++ 4 files changed, 811 insertions(+), 1 deletion(-) create mode 100644 tests/integration/_support/bedrock_runtime_peer.py create mode 100644 tests/integration/_support/provider.py create mode 100644 tests/integration/providers/_cache_control_marks_support.py diff --git a/tests/integration/_support/bedrock_runtime_peer.py b/tests/integration/_support/bedrock_runtime_peer.py new file mode 100644 index 00000000000..64115d9eaa0 --- /dev/null +++ b/tests/integration/_support/bedrock_runtime_peer.py @@ -0,0 +1,278 @@ +from __future__ import annotations + +import json +import re +import threading +from collections.abc import Mapping +from multiprocessing.sharedctypes import Synchronized +from types import MappingProxyType +from typing import Final +from urllib.parse import unquote + +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +REASONING_EFFORTS: Final = frozenset(("none", "minimal", "low", "medium", "high", "xhigh")) +NATIVE_CHAT: Final = "/openai/v1/chat/completions" +NATIVE_RESPONSES: Final = "/openai/v1/responses" +PNG_1X1: Final = bytes.fromhex( + "89504e470d0a1a0a0000000d49484452000000010000000108060000001f15c489" + "0000000d49444154789c63f8cfc0f01f00050001ff89993d1d0000000049454e44ae426082" +) +USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "prompt_tokens": 9, + "completion_tokens": 5, + "total_tokens": 14, + "completion_tokens_details": {"reasoning_tokens": 3}, + } +) +_STATUS: Final = re.compile(r"status=(\d{3})") +_CONVERSE: Final = re.compile(r"^/model/(.+)/converse$") +_CONVERSE_STREAM: Final = re.compile(r"^/model/(.+)/converse-stream$") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_NO_MARKER: Final = "0" * 32 + + +def marker_of(request: Request) -> str: + found: Final = MARKER.search(request.body.decode(errors="replace")) + return _NO_MARKER if found is None else found.group(1) + + +def body_of(request: Request) -> Mapping[str, JsonValue]: + try: + return _JSON_OBJECT.validate_json(request.body) + except ValueError: + return {} + + +def target_of(request: Request) -> str: + return unquote(request.target) + + +def answer(marker: str) -> str: + return f"answer marker-{marker}" + + +def reasoning_answer(marker: str) -> str: + return f"why marker-{marker} {answer(marker)}" + + +def _headers(marker: str) -> Mapping[str, str]: + return MappingProxyType({"x-amzn-requestid": marker}) + + +def _json_reply(status: int, payload: Mapping[str, JsonValue], marker: str) -> Reply: + return Reply(status=status, body=json.dumps(payload).encode(), headers=_headers(marker)) + + +def _error(status: int, message: str, marker: str) -> Reply: + return _json_reply(status, {"message": message}, marker) + + +def _effort_of(target: str, body: Mapping[str, JsonValue]) -> JsonValue: + if not _CONVERSE.match(target) and not _CONVERSE_STREAM.match(target): + return body.get("reasoning_effort") + fields: Final = body.get("additionalModelRequestFields") + reasoning: Final = fields.get("reasoning") if isinstance(fields, Mapping) else None + return reasoning.get("effort") if isinstance(reasoning, Mapping) else None + + +def forwarded_effort(request: Request) -> JsonValue: + return _effort_of(target_of(request), body_of(request)) + + +def _sse(frames: tuple[Mapping[str, JsonValue], ...], pause: float) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=(*(b"data: " + json.dumps(frame).encode() + b"\n\n" for frame in frames), b"data: [DONE]\n\n"), + pause_between_chunks=pause, + ) + + +def _with_headers(reply: Reply, marker: str) -> Reply: + return Reply( + status=reply.status, + body=reply.body, + content_type=reply.content_type, + chunks=reply.chunks, + abort_after=reply.abort_after, + gate_after_first=reply.gate_after_first, + pause_between_chunks=reply.pause_between_chunks, + headers=_headers(marker), + ) + + +def _content_deltas(model: str, marker: str) -> tuple[str, ...]: + if "gpt-oss" in model: + return ("why ", f"marker-{marker}", " answer ", f"marker-{marker}") + return ("answer ", f"marker-{marker}") + + +def _chat_text(model: str, marker: str) -> str: + return reasoning_answer(marker) if "gpt-oss" in model else answer(marker) + + +def _chat_reply(model: str, marker: str, stream: bool, pause: float) -> Reply: + identity: Final = f"chatcmpl-{marker}" + if not stream: + return _json_reply( + 200, + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": _chat_text(model, marker)}, + "finish_reason": "stop", + } + ], + "usage": dict(USAGE), + }, + marker, + ) + deltas: Final = _content_deltas(model, marker) + frames: Final = tuple( + { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": delta}, "finish_reason": None}], + } + for delta in deltas + ) + finish: Final[Mapping[str, JsonValue]] = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": dict(USAGE), + } + return _with_headers(_sse((*frames, finish), pause), marker) + + +def _responses_reply(model: str, marker: str, stream: bool, pause: float) -> Reply: + identity: Final = f"resp_upstream_{marker}" + item_id: Final = f"msg_{marker}" + response: Final[Mapping[str, JsonValue]] = { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": model, + "output": [ + { + "type": "message", + "id": item_id, + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": answer(marker), "annotations": []}], + } + ], + "usage": {"input_tokens": 30, "output_tokens": 5, "total_tokens": 35}, + } + if not stream: + return _json_reply(200, response, marker) + events: Final[tuple[Mapping[str, JsonValue], ...]] = ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**response, "status": "in_progress", "output": []}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": item_id, + "output_index": 0, + "content_index": 0, + "delta": answer(marker), + }, + {"type": "response.completed", "sequence_number": 2, "response": response}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events), + pause_between_chunks=pause, + headers=_headers(marker), + ) + + +def _converse_reply(marker: str) -> Reply: + return _json_reply( + 200, + { + "output": {"message": {"role": "assistant", "content": [{"text": answer(marker)}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 9, "outputTokens": 5, "totalTokens": 14}, + "metrics": {"latencyMs": 1}, + }, + marker, + ) + + +def _converse_stream_reply(marker: str, pause: float) -> Reply: + events: Final[tuple[tuple[str, Mapping[str, JsonValue]], ...]] = ( + ("messageStart", {"role": "assistant"}), + ("contentBlockDelta", {"delta": {"text": "answer "}, "contentBlockIndex": 0}), + ("contentBlockDelta", {"delta": {"text": f"marker-{marker}"}, "contentBlockIndex": 0}), + ("contentBlockStop", {"contentBlockIndex": 0}), + ("messageStop", {"stopReason": "end_turn"}), + ("metadata", {"usage": {"inputTokens": 9, "outputTokens": 5, "totalTokens": 14}, "metrics": {"latencyMs": 1}}), + ) + return Reply( + content_type=EVENT_STREAM, + chunks=tuple(_aws_event_frame(kind, payload, "sc", marker) for kind, payload in events), + pause_between_chunks=pause, + headers=_headers(marker), + ) + + +def respond(request: Request, *, pause: float = 0.0) -> Reply: + target: Final = target_of(request) + marker: Final = marker_of(request) + if request.method == "GET": + if target == "/image.png": + return Reply(body=PNG_1X1, content_type="image/png", headers=_headers(marker)) + return _error(404, f"no scripted object at {target}", marker) + scripted_status: Final = _STATUS.search(request.body.decode(errors="replace")) + if scripted_status is not None: + status: Final = int(scripted_status.group(1)) + return _error(status, f"scripted {status}", marker) + body: Final = body_of(request) + effort: Final = _effort_of(target, body) + if effort is not None and (not isinstance(effort, str) or effort not in REASONING_EFFORTS): + return _error(400, f"Invalid reasoning effort: {json.dumps(effort)}", marker) + model: Final = str(body.get("model", "")) + stream: Final = body.get("stream") is True + if request.method == "POST" and target == NATIVE_CHAT: + return _chat_reply(model, marker, stream, pause) + if request.method == "POST" and target == NATIVE_RESPONSES: + return _responses_reply(model, marker, stream, pause) + if request.method == "POST" and _CONVERSE.match(target): + return _converse_reply(marker) + if request.method == "POST" and _CONVERSE_STREAM.match(target): + return _converse_stream_reply(marker, pause) + return _error(404, f"unknown bedrock route {request.method} {target}", marker) + + +def serve_peer(port: int, received: Synchronized[int], answer_first: int) -> None: + held: Final = threading.Event() + + def respond_or_hold(request: Request) -> Reply: + with received.get_lock(): + received.value += 1 + ordinal: Final = received.value + if ordinal > answer_first: + held.wait() + return respond(request) + + with wire_server(respond_or_hold, port=port): + threading.Event().wait() diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index fcbaf7c8d8c..beb23a611e6 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -46,12 +46,16 @@ def signal_group(group: int, action: int) -> None: pass +def graceful_stop_seconds() -> float: + return 30.0 + + def stop_root_process(process: subprocess.Popen[bytes]) -> bool: if process.poll() is not None: return True process.terminate() try: - process.wait(timeout=30) + process.wait(timeout=graceful_stop_seconds()) except subprocess.TimeoutExpired: return False return True diff --git a/tests/integration/_support/provider.py b/tests/integration/_support/provider.py new file mode 100644 index 00000000000..fa0bb7ad379 --- /dev/null +++ b/tests/integration/_support/provider.py @@ -0,0 +1,49 @@ +"""The fake provider shared by every integration test that takes the `provider` fixture. + +Deployments in `proxy_config.yaml` point at `PROVIDER_URL`, so one server answers for all of them. A test +queues the replies it expects with `expect` and reads what the proxy sent with `received`. Tests run one at a +time against it; the `provider` fixture checks nothing is left over between tests. +The proxy's model-info refresh is answered here. +""" + +from __future__ import annotations + +import os +from collections import deque +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Final + +from tests.integration._support.wire import Reply, Request, Wire, wire_server + +PROVIDER_PORT: Final = int(os.environ.get("INTEGRATION_PROVIDER_PORT", "8191")) +PROVIDER_URL: Final = f"http://127.0.0.1:{PROVIDER_PORT}" +MODEL_DISCOVERY: Final = ("GET", "/v1/models") +_UNQUEUED: Final = Reply(status=500, body=b'{"error": "the shared fake provider has no reply queued for this request"}') + + +@dataclass(slots=True) +class SharedProvider: + wire: Wire + replies: deque[Reply] + last_test: str | None = field(default=None) + + def expect(self, *replies: Reply) -> None: + self.replies.extend(replies) + + def received(self) -> tuple[Request, ...]: + return tuple(request for request in self.wire.drain() if (request.method, request.target) != MODEL_DISCOVERY) + + +@contextmanager +def shared_provider() -> Iterator[SharedProvider]: + replies: Final[deque[Reply]] = deque() + + def respond(request: Request) -> Reply: + if (request.method, request.target) == MODEL_DISCOVERY: + return Reply(body=b'{"object":"list","data":[]}') + return replies.popleft() if replies else _UNQUEUED + + with wire_server(respond, port=PROVIDER_PORT) as wire: + yield SharedProvider(wire, replies) diff --git a/tests/integration/providers/_cache_control_marks_support.py b/tests/integration/providers/_cache_control_marks_support.py new file mode 100644 index 00000000000..2f6071f9e45 --- /dev/null +++ b/tests/integration/providers/_cache_control_marks_support.py @@ -0,0 +1,479 @@ +import json +import re +import uuid +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +# https://platform.claude.com/docs/en/build-with-claude/prompt-caching (read 2026-09-28): at most 4 blocks with cache_control +ANTHROPIC_CACHE_CONTROL_CAP: Final = 4 +# https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html (read 2026-09-28): at most 4 cache checkpoints +BEDROCK_CACHE_CHECKPOINT_CAP: Final = 4 + +ANTHROPIC_MODEL: Final = "claude-opus-5-5" +BEDROCK_MODEL: Final = "anthropic.claude-opus-5-5" +UNMAPPED_BEDROCK_ARN: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/scripted" +PROVIDER_KEY: Final = "synthetic-provider-key" +SYSTEM: Final = "Use the provided tool results to answer the user." +ASK: Final = "Look up the weather in London, Paris, and Tokyo." +CITIES: Final = ("London", "Paris", "Tokyo") +EPHEMERAL: Final[dict[str, JsonValue]] = {"type": "ephemeral"} +POINTS: Final[list[JsonValue]] = [ + {"location": "message", "role": "system"}, + {"location": "message", "index": -1}, +] +TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": { + "name": "lookup_weather", + "description": "Look up the weather in a city.", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} +SYSTEM_LABEL: Final = f"system:{SYSTEM}" +ASK_LABEL: Final = f"user:text:{ASK}" +_MARKER: Final = re.compile(rb"marker-([0-9a-f]{32})") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class Mark: + label: str + ttl: str | None + + +def new_marker() -> str: + return uuid.uuid4().hex + + +def final_text(marker: str) -> str: + return f"Summarize the results in one word. marker-{marker}" + + +def final_label(marker: str) -> str: + return f"user:text:{final_text(marker)}" + + +def call_id(city: str) -> str: + return f"call_weather_{city.lower()}" + + +def tool_use_label(city: str) -> str: + return f"assistant:tool_use:{call_id(city)}" + + +def tool_call(city: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "id": call_id(city), + "type": "function", + "function": {"name": "lookup_weather", "arguments": json.dumps({"city": city})}, + **fields, + } + + +def marked_calls(mark: JsonValue = EPHEMERAL, cities: Sequence[str] = CITIES) -> list[JsonValue]: + return [tool_call(city, cache_control=mark) for city in cities] + + +def client_marked() -> list[str]: + return [ASK_LABEL, *(tool_use_label(city) for city in CITIES)] + + +def ask(*, marked: bool) -> dict[str, JsonValue]: + block: Final[dict[str, JsonValue]] = {"type": "text", "text": ASK} + return {"role": "user", "content": [{**block, "cache_control": EPHEMERAL} if marked else block]} + + +def tool_results(calls: Sequence[JsonValue]) -> list[JsonValue]: + return [ + {"role": "tool", "tool_call_id": call["id"], "content": "sunny."} + for call in calls + if isinstance(call, dict) and str(call.get("id", "")).startswith("call_") + ] + + +def conversation( + marker: str, + calls: Sequence[JsonValue], + *, + ask_marked: bool = True, + assistant: dict[str, JsonValue] | None = None, +) -> list[JsonValue]: + turn: Final[dict[str, JsonValue]] = { + "role": "assistant", + "content": "", + "tool_calls": list(calls), + **(assistant or {}), + } + return [ + {"role": "system", "content": SYSTEM}, + ask(marked=ask_marked), + turn, + *tool_results(calls), + {"role": "user", "content": final_text(marker)}, + ] + + +def chat_body(model: str, messages: Sequence[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "messages": list(messages), "tools": [TOOL], "max_tokens": 64, **fields} + + +def messages_body(model: str, marker: str, *, stream: bool) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 64, + "stream": stream, + "system": [{"type": "text", "text": SYSTEM}], + "tools": [ + { + "name": "lookup_weather", + "description": "Look up the weather in a city.", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + ], + "messages": [ + {"role": "user", "content": [{"type": "text", "text": ASK, "cache_control": EPHEMERAL}]}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": call_id(city), + "name": "lookup_weather", + "input": {"city": city}, + "cache_control": EPHEMERAL, + } + for city in CITIES + ], + }, + { + "role": "user", + "content": [ + *({"type": "tool_result", "tool_use_id": call_id(city), "content": "sunny."} for city in CITIES), + {"type": "text", "text": final_text(marker)}, + ], + }, + ], + } + + +def responses_body(model: str, marker: str) -> dict[str, JsonValue]: + items: Final[list[JsonValue]] = [ + {"role": "user", "content": [{"type": "input_text", "text": ASK, "cache_control": EPHEMERAL}]}, + *( + { + "type": "function_call", + "call_id": call_id(city), + "name": "lookup_weather", + "arguments": json.dumps({"city": city}), + "cache_control": EPHEMERAL, + } + for city in CITIES + ), + *({"type": "function_call_output", "call_id": call_id(city), "output": "sunny."} for city in CITIES), + {"role": "user", "content": final_text(marker)}, + ] + return { + "model": model, + "instructions": SYSTEM, + "max_output_tokens": 64, + "tools": [{"type": "function", "name": "lookup_weather", "parameters": {"type": "object"}}], + "input": items, + } + + +def marker_of(request: Request) -> str: + found: Final = _MARKER.findall(request.body) + return found[-1].decode() if found else "unmarked" + + +class _AnthropicBlock(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + type: str + text: str | None = None + id: str | None = None + tool_use_id: str | None = None + name: str | None = None + cache_control: JsonValue = None + + +class _AnthropicMessage(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + role: str + content: str | tuple[_AnthropicBlock, ...] + + +class _AnthropicTool(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + name: str = "" + cache_control: JsonValue = None + + +class _AnthropicBody(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + system: str | tuple[_AnthropicBlock, ...] = () + messages: tuple[_AnthropicMessage, ...] = () + tools: tuple[_AnthropicTool, ...] = () + cache_control: JsonValue = None + stream: bool = False + + +def _ttl(cache_control: JsonValue) -> str | None: + if not isinstance(cache_control, dict): + return None + ttl: Final = cache_control.get("ttl") + return ttl if isinstance(ttl, str) else None + + +def _block_label(role: str, block: _AnthropicBlock) -> str: + detail: Final = block.text if block.type == "text" else block.id or block.tool_use_id or "" + return f"{role}:{block.type}:{detail}" + + +def _message_blocks(body: _AnthropicBody) -> Iterator[tuple[str, _AnthropicBlock]]: + for message in body.messages: + if isinstance(message.content, tuple): + yield from ((message.role, block) for block in message.content) + + +def _anthropic_body(request: Request) -> _AnthropicBody: + return _AnthropicBody.model_validate_json(request.body) + + +def anthropic_marks(request: Request) -> tuple[Mark, ...]: + body: Final = _anthropic_body(request) + system: Final = body.system if isinstance(body.system, tuple) else () + return ( + *(Mark(f"tool:{tool.name}", _ttl(tool.cache_control)) for tool in body.tools if tool.cache_control is not None), + *( + Mark(f"system:{block.text}", _ttl(block.cache_control)) + for block in system + if block.cache_control is not None + ), + *( + Mark(_block_label(role, block), _ttl(block.cache_control)) + for role, block in _message_blocks(body) + if block.cache_control is not None + ), + *((Mark("request", _ttl(body.cache_control)),) if body.cache_control is not None else ()), + ) + + +def anthropic_labels(request: Request) -> list[str]: + return [mark.label for mark in anthropic_marks(request)] + + +def _anthropic_error(message: str) -> Reply: + return Reply( + status=400, + body=json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": message}}).encode(), + ) + + +def _ttl_out_of_order(marks: Sequence[Mark]) -> bool: + first_short: Final = next((index for index, mark in enumerate(marks) if mark.ttl != "1h"), len(marks)) + return any(mark.ttl == "1h" for mark in marks[first_short:]) + + +def _usage() -> dict[str, JsonValue]: + return {"input_tokens": 12, "output_tokens": 1, "cache_creation_input_tokens": 12, "cache_read_input_tokens": 0} + + +def _anthropic_message(identity: str) -> dict[str, JsonValue]: + return { + "id": identity, + "type": "message", + "role": "assistant", + "model": ANTHROPIC_MODEL, + "content": [{"type": "text", "text": "sunny"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": _usage(), + } + + +def _event(name: str, data: dict[str, JsonValue]) -> bytes: + return f"event: {name}\ndata: {json.dumps(data)}\n\n".encode() + + +def _anthropic_stream(identity: str) -> tuple[bytes, ...]: + return ( + _event( + "message_start", + {"type": "message_start", "message": {**_anthropic_message(identity), "content": [], "stop_reason": None}}, + ), + _event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "sunny"}}, + ), + _event("content_block_stop", {"type": "content_block_stop", "index": 0}), + _event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 1}, + }, + ), + _event("message_stop", {"type": "message_stop"}), + ) + + +def anthropic_peer(request: Request) -> Reply: + marks: Final = anthropic_marks(request) + if len(marks) > ANTHROPIC_CACHE_CONTROL_CAP: + return _anthropic_error( + f"A maximum of {ANTHROPIC_CACHE_CONTROL_CAP} blocks with cache_control may be provided. Found {len(marks)}." + ) + if _ttl_out_of_order(marks): + return _anthropic_error("a ttl='1h' cache_control block must not come after a ttl='5m' cache_control block") + identity: Final = f"msg_{marker_of(request)}" + if _anthropic_body(request).stream: + return Reply(content_type="text/event-stream", chunks=_anthropic_stream(identity)) + return Reply(body=json.dumps(_anthropic_message(identity)).encode()) + + +def _bedrock_label(role: str, block: JsonValue) -> str: + if not isinstance(block, dict): + return f"{role}:start" + if isinstance(block.get("text"), str): + return f"{role}:text:{block['text']}" + for kind in ("toolUse", "toolResult"): + inner = block.get(kind) + if isinstance(inner, dict): + return f"{role}:{kind}:{inner.get('toolUseId')}" + spec: Final = block.get("toolSpec") + return f"tool:{spec.get('name')}" if isinstance(spec, dict) else f"{role}:other" + + +def _cache_points(role: str, blocks: JsonValue) -> Iterator[str]: + listed: Final = blocks if isinstance(blocks, list) else [] + for previous, block in zip([None, *listed], listed): + if isinstance(block, dict) and "cachePoint" in block: + yield _bedrock_label(role, previous) + + +def _bedrock_sections(body: dict[str, JsonValue]) -> Iterator[tuple[str, JsonValue]]: + tool_config: Final = body.get("toolConfig") + yield ("tool", tool_config.get("tools") if isinstance(tool_config, dict) else None) + yield ("system", body.get("system")) + messages: Final = body.get("messages") + for message in messages if isinstance(messages, list) else []: + if isinstance(message, dict): + yield (str(message.get("role")), message.get("content")) + + +def bedrock_labels(request: Request) -> list[str]: + body: Final = _JSON_OBJECT.validate_json(request.body) + return [label for role, blocks in _bedrock_sections(body) for label in _cache_points(role, blocks)] + + +def bedrock_peer(request: Request) -> Reply: + found: Final = len(bedrock_labels(request)) + if found > BEDROCK_CACHE_CHECKPOINT_CAP: + return Reply( + status=400, + headers={"x-amzn-errortype": "ValidationException"}, + body=json.dumps( + { + "message": f"A maximum of {BEDROCK_CACHE_CHECKPOINT_CAP} cache checkpoints may be provided. Found {found}." + } + ).encode(), + ) + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "sunny"}]}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 12, + "outputTokens": 1, + "totalTokens": 13, + "cacheWriteInputTokens": 12, + "cacheReadInputTokens": 0, + }, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + +def _without_cache_points(value: JsonValue) -> JsonValue: + if isinstance(value, dict): + return {key: _without_cache_points(inner) for key, inner in value.items()} + if isinstance(value, list): + return [ + _without_cache_points(inner) for inner in value if not (isinstance(inner, dict) and "cachePoint" in inner) + ] + return value + + +def sent_without_cache_points(requests: Sequence[Request], marker: str) -> list[JsonValue]: + return [ + _without_cache_points(_JSON_OBJECT.validate_json(request.body.replace(marker.encode(), b"marker"))) + for request in requests + ] + + +def gateway_injected(response_id: str) -> bool: + rows: Final = eventually( + lambda: read_rows( + """SELECT metadata->>'litellm_gateway_injected_cache' AS injected FROM "LiteLLM_SpendLogs" """ + "WHERE request_id=%s", + (response_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return rows[0]["injected"] is not None + + +def post_chat(gateway: Gateway, body: dict[str, JsonValue], *, key: str | None = None) -> tuple[int, str, str]: + response: Final = gateway.request("POST", "/v1/chat/completions", body, key=key) + identity: Final = _JSON_OBJECT.validate_json(response.content).get("id") if response.status_code == 200 else None + return response.status_code, str(identity), response.text + + +def anthropic_deployment(name: str, api_base: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": f"anthropic/{ANTHROPIC_MODEL}", + "api_base": api_base, + "api_key": PROVIDER_KEY, + **fields, + }, + } + + +def owned_config( + directory: Path, + model_list: Sequence[JsonValue], + *, + litellm_settings: Mapping[str, JsonValue] = MappingProxyType({}), + router_settings: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Path: + config: Final = _JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + merged: Final = { + **config, + "model_list": list(model_list), + "litellm_settings": {**object_value(config["litellm_settings"]), **litellm_settings}, + "router_settings": {**object_value(config["router_settings"]), "num_retries": 0, **router_settings}, + } + path: Final = directory / f"cache-control-marks-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(merged)) + return path From b6f1a033deda54461feda0a47fa59532aee5c089 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:14:29 -0700 Subject: [PATCH 23/36] feat(count_tokens)!: configurable token_count_framing and token_count_model, count source in the response (#46777) Backport of #46777 to stable/1.104.x. Cherry-picked from 75faaeab641e161a3d6bc38a42fcaa71c5602d90 (main). --- litellm/litellm_core_utils/core_helpers.py | 26 +- .../prompt_templates/factory.py | 22 + litellm/llms/bedrock/count_tokens/kimi_k3.py | 502 ++++++ .../responses/count_tokens/token_counter.py | 1 + .../llms/vertex_ai/count_tokens/context.py | 25 +- ...odel_prices_and_context_window_backup.json | 28 +- .../proxy/anthropic_endpoints/endpoints.py | 3 +- .../common_utils/token_count_settings.py | 164 ++ litellm/proxy/google_endpoints/endpoints.py | 11 +- litellm/proxy/proxy_server.py | 461 +++++- .../proxy/response_api_endpoints/endpoints.py | 5 +- litellm/types/llms/vertex_ai.py | 6 +- litellm/types/token_count.py | 16 + litellm/types/utils.py | 9 + litellm/utils.py | 3 + model_prices_and_context_window.json | 28 +- model_prices_and_context_window.schema.json | 30 + ...anthropic_count_tokens_system_lift_wire.py | 24 +- .../test_bedrock_mantle_count_tokens_wire.py | 89 +- .../test_vertex_partner_count_tokens_wire.py | 28 +- .../integration/spend/test_cache_and_quota.py | 8 +- .../test_google_api_endpoints.py | 42 +- .../test_model_management_endpoints.py | 48 + .../proxy/proxy_server/test_routes_utils.py | 1 + .../count_tokens/test_token_counter.py | 5 + .../count_tokens/test_bedrock_kimi_k3.py | 524 +++++++ .../count_tokens/test_token_counter.py | 5 + .../common_utils/test_token_count_settings.py | 190 +++ tests/unit/proxy/test_proxy_token_counter.py | 1364 ++++++++++++++++- .../unit/test_router_model_cost_isolation.py | 34 + tests/unit/test_utils.py | 7 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 +- 32 files changed, 3571 insertions(+), 150 deletions(-) create mode 100644 litellm/llms/bedrock/count_tokens/kimi_k3.py create mode 100644 litellm/proxy/common_utils/token_count_settings.py create mode 100644 litellm/types/token_count.py create mode 100644 tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py create mode 100644 tests/unit/proxy/common_utils/test_token_count_settings.py diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index b095b4b12c6..d74007dfccb 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -3,9 +3,9 @@ import copy import logging import re -from collections.abc import Collection, Iterable, Mapping +from collections.abc import Collection, Iterable, Iterator, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, runtime_checkable import httpx from pydantic import TypeAdapter, ValidationError @@ -51,6 +51,28 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non _DROP_PARAMS_BOOL: Final = TypeAdapter(bool) +def as_str_mapping(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, Mapping) else None + + +@runtime_checkable +class _SupportsIter(Protocol): + def __iter__(self) -> Iterator[object]: ... + + +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, index: int, /) -> object: ... + + +def iter_items(value: object) -> Iterable[object]: + if isinstance(value, _SupportsIter): + return value + if isinstance(value, _SupportsGetItem): + return iter(value) + raise TypeError(f"'{type(value).__name__}' object is not iterable") + + def normalize_drop_params(value: object) -> bool | None: if value is None or isinstance(value, bool): return value diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 7b12d1e939f..26d219a5ebc 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5580,3 +5580,25 @@ def resolve_structured_messages( if structured: return [msg if isinstance(msg, dict) else msg.model_dump() for msg in structured] return None + + +def is_flat_responses_tool(tool: Mapping[str, object]) -> bool: + tool_type: Final = tool.get("type") + return tool_type in ("function", "custom") and tool_type not in tool and "input_schema" not in tool + + +def chat_completion_tools_for(tool: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + if is_flat_responses_tool(tool): + chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + [dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in a flat Responses tool shape + ) + return tuple(chat_tools) + anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai( + [dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in an Anthropic or chat tool shape + ) + return tuple(anthropic_tools) diff --git a/litellm/llms/bedrock/count_tokens/kimi_k3.py b/litellm/llms/bedrock/count_tokens/kimi_k3.py new file mode 100644 index 00000000000..04ae9fae9c2 --- /dev/null +++ b/litellm/llms/bedrock/count_tokens/kimi_k3.py @@ -0,0 +1,502 @@ +""" +Exact Bedrock token counts for Kimi K3, which has no CountTokens support. + +Renders the Converse body litellm sends through Moonshot's Kimi K3 chat template and counts it with the +official tokenizer, downloaded once from Hugging Face at a pinned revision and checked against its SHA-256. +""" + +import asyncio +import base64 +import copy +import hashlib +import json +import threading +import time +from collections.abc import Callable, Iterable, Mapping, Sequence +from concurrent.futures import Future +from dataclasses import dataclass +from enum import Enum +from itertools import chain, groupby +from pathlib import Path +from typing import Annotated, Final, Literal, TypeAlias + +import regex +import tiktoken +from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError, model_validator +from typing_extensions import Self + +import litellm +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.prompt_templates.factory import ( + chat_completion_tools_for, + is_flat_responses_tool, +) +from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam + +KIMI_K3_REPO: Final = "moonshotai/Kimi-K3" +KIMI_K3_REVISION: Final = "f831ab66814297da540d832a5235f8e904f29d06" +KIMI_K3_RANKS_FILE: Final = "tiktoken.model" +KIMI_K3_RANKS_SHA256: Final = "b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103" +_PATTERN: Final = "|".join( + ( + r"[\p{Han}]+", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?", + r"\p{N}{1,3}", + r" ?[^\s\p{L}\p{N}]+[\r\n]*", + r"\s*[\r\n]+", + r"\s+(?!\S)", + r"\s+", + ) +) +# Ends chunks only where the pattern above always splits, so their counts sum to the whole text's count while BPE +# stays linear. Only a run of over 1024 characters with no such split, like 2000 spaces, is cut mid-piece. +_CHUNK: Final = regex.compile( + r"(?s).{1,1024}(?:(?<=\S)(?= \S)|(?<=\p{N})(?=\P{N})|(?<=[^\s\p{N}])(?=\p{N})|(?<=\p{L})(?=[^\p{L}\p{M}'])" + r"|(?<=[\r\n])(?=\S)|\Z)|.{1024}\p{N}*" +) + + +@dataclass(frozen=True, slots=True) +class KimiK3TokenizerSource: + ranks_path: Callable[[], str] + sha256: str + + +def _hugging_face_ranks_path() -> str: + import huggingface_hub + + return huggingface_hub.hf_hub_download( + repo_id=KIMI_K3_REPO, + filename=KIMI_K3_RANKS_FILE, + revision=KIMI_K3_REVISION, + local_files_only=litellm.disable_hf_tokenizer_download is True, + ) + + +HUGGING_FACE_SOURCE: Final = KimiK3TokenizerSource(ranks_path=_hugging_face_ranks_path, sha256=KIMI_K3_RANKS_SHA256) + + +def _loaded(source: KimiK3TokenizerSource) -> tiktoken.Encoding | str: + try: + data: Final = Path(source.ranks_path()).read_bytes() + except Exception as e: # noqa: BLE001 # a download or read failure of any kind falls back to the estimate + return type(e).__name__ + if hashlib.sha256(data).hexdigest() != source.sha256: + return "SHA-256 mismatch" + rows: Final = (line.split() for line in data.splitlines() if line) + return tiktoken.Encoding( + name="kimi_k3", + pat_str=_PATTERN, + mergeable_ranks={base64.b64decode(token): int(rank) for token, rank in rows}, + special_tokens={}, + ) + + +def load_kimi_k3_encoding(source: KimiK3TokenizerSource) -> tiktoken.Encoding | None: + loaded: Final = _loaded(source) + if isinstance(loaded, tiktoken.Encoding): + return loaded + verbose_logger.warning( + "Kimi K3 tokenizer unavailable (%s), so Kimi K3 token counts are estimated until a later count loads it. " + "Allow downloads of %s from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache", + loaded, + KIMI_K3_RANKS_FILE, + KIMI_K3_REPO, + KIMI_K3_REVISION, + ) + return None + + +def _started_load(source: KimiK3TokenizerSource) -> "Future[tiktoken.Encoding | None]": + encoding: Final = Future[tiktoken.Encoding | None]() + threading.Thread( + target=lambda: encoding.set_result(load_kimi_k3_encoding(source)), name="litellm-kimi-k3-tokenizer", daemon=True + ).start() + return encoding + + +@dataclass(frozen=True, slots=True) +class _EncodingLoad: + encoding: "Future[tiktoken.Encoding | None]" + started_at: float + is_retry: bool + + def retry_due(self, now: float, retry_after_seconds: float) -> bool: + return self.encoding.done() and self.encoding.result() is None and now - self.started_at >= retry_after_seconds + + +@dataclass(slots=True) +class KimiK3Encodings: + source: KimiK3TokenizerSource + retry_after_seconds: float + clock: Callable[[], float] = time.monotonic + current: _EncodingLoad | None = None + + def _load(self) -> _EncodingLoad: + now: Final = self.clock() + current: Final = self.current + if current is not None and not current.retry_due(now, self.retry_after_seconds): + return current + load: Final = _EncodingLoad(encoding=_started_load(self.source), started_at=now, is_retry=current is not None) + self.current = load + return load + + async def encoding(self, wait_seconds: float) -> tiktoken.Encoding | None: + load: Final = self._load() + remaining: Final = 0 if load.is_retry else load.started_at + wait_seconds - self.clock() + if load.encoding.done() or remaining <= 0: + return load.encoding.result() if load.encoding.done() else None + await asyncio.wait((asyncio.wrap_future(load.encoding),), timeout=remaining) + if load.encoding.done(): + return load.encoding.result() + verbose_logger.warning( + "Kimi K3 tokenizer still loading after %s s, so Kimi K3 token counts are estimated until it loads. " + "Check that %s can be downloaded from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache", + wait_seconds, + KIMI_K3_RANKS_FILE, + KIMI_K3_REPO, + KIMI_K3_REVISION, + ) + return None + + +class _Converse(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class _Text(_Converse): + text: str + + +class _CachePoint(_Converse): + cache_point: JsonValue = Field(alias="cachePoint") + + +class _ToolUse(_Converse): + tool_use_id: str = Field(alias="toolUseId") + name: str + arguments: Mapping[str, JsonValue] = Field(alias="input") + + +class _ToolUseBlock(_Converse): + tool_use: _ToolUse = Field(alias="toolUse") + + +class _ToolResult(_Converse): + tool_use_id: str = Field(alias="toolUseId") + content: tuple[_Text, ...] + status: str | None = None + + +class _ToolResultBlock(_Converse): + tool_result: _ToolResult = Field(alias="toolResult") + + +class _ReasoningText(_Converse): + text: str + signature: str | None = None + + +class _Reasoning(_Converse): + reasoning_text: _ReasoningText = Field(alias="reasoningText") + + +class _ReasoningBlock(_Converse): + reasoning_content: _Reasoning = Field(alias="reasoningContent") + + +class _UserMessage(_Converse): + role: Literal["user"] + content: tuple[_Text | _ToolResultBlock | _CachePoint, ...] + + +class _AssistantMessage(_Converse): + role: Literal["assistant"] + content: tuple[_Text | _ToolUseBlock | _ReasoningBlock | _CachePoint, ...] + + def calls(self) -> tuple[_ToolUse, ...]: + return tuple(block.tool_use for block in self.content if isinstance(block, _ToolUseBlock)) + + +class _InputSchema(_Converse): + definition: JsonValue = Field(alias="json") + + +class _ToolSpec(_Converse): + name: str + description: str | None = None + input_schema: _InputSchema = Field(alias="inputSchema") + + def declared(self) -> Mapping[str, JsonValue]: + described: Final[Mapping[str, JsonValue]] = ( + {} if self.description is None else {"description": self.description} + ) + return {"name": self.name, **described, "parameters": self.input_schema.definition} + + +class _ToolSpecBlock(_Converse): + tool_spec: _ToolSpec = Field(alias="toolSpec") + + +class _ToolConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + tools: tuple[_ToolSpecBlock | _CachePoint, ...] + + +_Message: TypeAlias = Annotated[_UserMessage | _AssistantMessage, Field(discriminator="role")] + + +def _result_ids(message: _Message) -> frozenset[str]: + if isinstance(message, _AssistantMessage): + return frozenset() + return frozenset(block.tool_result.tool_use_id for block in message.content if isinstance(block, _ToolResultBlock)) + + +def _call_ids(message: _Message | None) -> frozenset[str]: + if not isinstance(message, _AssistantMessage): + return frozenset() + return frozenset(call.tool_use_id for call in message.calls()) + + +class _ConverseRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + messages: tuple[_Message, ...] + system: tuple[_Text | _CachePoint, ...] = () + tool_config: _ToolConfig | None = Field(default=None, alias="toolConfig") + + @model_validator(mode="after") + def _results_answer_the_previous_calls(self) -> Self: + pairs: Final = zip((None, *self.messages), self.messages, strict=False) + if not all(_result_ids(message) <= _call_ids(previous) for previous, message in pairs): + raise ValueError("a toolResult does not answer a toolUse of the previous assistant message") + return self + + +class _Control(Enum): + OPEN = "<|open|>" + CLOSE = "<|close|>" + SEP = "<|sep|>" + END_OF_MSG = "<|end_of_msg|>" + + +_Segment: TypeAlias = str | _Control +_Segments: TypeAlias = tuple[_Segment, ...] + + +def _open(tag: str, attrs: Iterable[tuple[str, str]] = ()) -> _Segments: + escaped: Final = ( + (f" {key}", '="', value.replace("&", "&").replace('"', """), '"') for key, value in attrs + ) + return (_Control.OPEN, tag, *chain.from_iterable(escaped), _Control.SEP) + + +def _close(tag: str) -> _Segments: + return (_Control.CLOSE, tag, _Control.SEP) + + +def _message(attrs: Iterable[tuple[str, str]], body: Iterable[_Segment]) -> _Segments: + return (*_open("message", attrs), *body, *_close("message"), _Control.END_OF_MSG) + + +def _tool_declare(specs: Sequence[_ToolSpec]) -> _Segments: + tools: Final = json.dumps( + [{"type": "function", "function": spec.declared()} for spec in specs], + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + return _message( + (("role", "system"), ("type", "tool-declare")), + (f"# Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n{tools}\n```",), + ) + + +def _text(blocks: Iterable[_Converse]) -> str: + return "".join(block.text for block in blocks if isinstance(block, _Text)) + + +def _json_type(value: JsonValue) -> str: + match value: + case bool(): + return "boolean" + case None: + return "null" + case int() | float(): + return "number" + case str(): + return "string" + case dict(): + return "object" + case _: + return "array" + + +def _argument(key: str, value: JsonValue) -> _Segments: + text: Final = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False) + return (*_open("argument", (("key", key), ("type", _json_type(value)))), text, *_close("argument")) + + +def _call(index: int, call: _ToolUse) -> _Segments: + arguments: Final = chain.from_iterable(_argument(key, value) for key, value in call.arguments.items()) + return (*_open("call", (("tool", call.name), ("index", str(index)))), *arguments, *_close("call")) + + +def _tool_calls(calls: Sequence[_ToolUse]) -> _Segments: + if not calls: + return () + rendered: Final = chain.from_iterable(_call(index, call) for index, call in enumerate(calls, start=1)) + return (*_open("tools"), *rendered, *_close("tools")) + + +def _assistant_message(message: _AssistantMessage) -> _Segments: + reasoning: Final = "".join( + block.reasoning_content.reasoning_text.text for block in message.content if isinstance(block, _ReasoningBlock) + ) + return _message( + (("role", "assistant"),), + ( + *_open("think"), + reasoning, + *_close("think"), + *_open("response"), + _text(message.content), + *_close("response"), + *_tool_calls(message.calls()), + ), + ) + + +def _tool_messages(results: Sequence[_ToolResult], previous: _AssistantMessage) -> _Segments: + names: Final = {call.tool_use_id: call.name for call in previous.calls()} + return tuple( + chain.from_iterable( + _message( + (("role", "tool"), ("tool", names[result.tool_use_id]), ("index", str(index))), + (_text(result.content),), + ) + for index, result in enumerate(results, start=1) + ) + ) + + +def _user_run(blocks: Sequence[_Text | _ToolResultBlock], previous: _Message | None) -> _Segments: + results: Final = tuple(block.tool_result for block in blocks if isinstance(block, _ToolResultBlock)) + if results and isinstance(previous, _AssistantMessage): + return _tool_messages(results, previous) + return _message((("role", "user"),), (_text(blocks),)) + + +def _user_messages(message: _UserMessage, previous: _Message | None) -> _Segments: + blocks: Final = tuple(block for block in message.content if not isinstance(block, _CachePoint)) + runs: Final = groupby(blocks, key=lambda block: isinstance(block, _ToolResultBlock)) + return tuple(chain.from_iterable(_user_run(tuple(run), previous) for _, run in runs)) + + +def _turn(previous: _Message | None, message: _Message) -> _Segments: + if isinstance(message, _AssistantMessage): + return _assistant_message(message) + return _user_messages(message, previous) + + +def _rendered(request: _ConverseRequest) -> _Segments: + specs: Final = ( + () + if request.tool_config is None + else tuple(block.tool_spec for block in request.tool_config.tools if isinstance(block, _ToolSpecBlock)) + ) + system: Final = _text(request.system) + turns: Final = chain.from_iterable( + _turn(previous, message) for previous, message in zip((None, *request.messages), request.messages, strict=False) + ) + return ( + *(_tool_declare(specs) if specs else ()), + *(_message((("role", "system"),), (system,)) if system else ()), + *turns, + *_open("message", (("role", "assistant"),)), + *_open("think"), + ) + + +def _text_tokens(encoding: tiktoken.Encoding, text: str) -> int: + spans: Final = (match.span() for match in _CHUNK.finditer(text)) + return sum(len(encoding.encode_ordinary(text[start:end])) for start, end in spans) + + +class _TextPart(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + type: Literal["text"] + + +class _TextOnlyMessage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + content: str | tuple[_TextPart, ...] | None = None + + +def _text_only(message: AllMessageValues) -> bool: + try: + _TextOnlyMessage.model_validate(message) + except ValidationError: + return False + return True + + +def _as_converse_input(tool: ChatCompletionToolParam) -> tuple[Mapping[str, object], ...]: + return chat_completion_tools_for(tool) if is_flat_responses_tool(tool) else (tool,) + + +def _converse_request( + model: str, messages: Sequence[AllMessageValues], tools: Sequence[ChatCompletionToolParam] | None +) -> _ConverseRequest | None: + if not all(_text_only(message) for message in messages): + return None + try: + chat_tools: Final = tuple( + chat_tool for tool in tools or () for chat_tool in _as_converse_input(tool) + ) # comprehension-ok: flattens each tool's translation + return _ConverseRequest.model_validate( + litellm.AmazonConverseConfig().transform_request( + model=model, + messages=copy.deepcopy(list(messages)), + optional_params={"tools": copy.deepcopy(list(chat_tools))} if chat_tools else {}, + litellm_params={}, + headers={}, + ) + ) + except Exception: # noqa: BLE001 # a body litellm cannot send to Bedrock, or cannot render, falls back to the estimate + return None + + +def kimi_k3_tokens( + encoding: tiktoken.Encoding, + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam] | None, +) -> int | None: + request: Final = _converse_request(model, messages, tools) + if request is None: + return None + return sum( + 1 if isinstance(segment, _Control) else _text_tokens(encoding, segment) for segment in _rendered(request) + ) + + +def kimi_k3_anthropic_tokens(encoding: tiktoken.Encoding, model: str, fields: Mapping[str, object]) -> int | None: + body: Final = {**fields, "model": model, "max_tokens": 1} + try: + translated, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + copy.deepcopy(body), # pyright: ignore[reportArgumentType] # raw /v1/messages fields + custom_llm_provider="bedrock", + ) + except Exception: # noqa: BLE001 # a request the adapter cannot translate falls back to the estimate + return None + match translated: + case {"tools": tools}: + return kimi_k3_tokens(encoding, model, translated["messages"], tools) + case _: + return kimi_k3_tokens(encoding, model, translated["messages"], None) + + +def kimi_k3_template_tokens(message_count: int, has_tools: bool) -> int: + return 14 + 13 * max(message_count - 1, 0) + (40 if has_tools else 0) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index bdaea0b4104..b34b6fbce2f 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -116,6 +116,7 @@ class OpenAITokenCounter(BaseTokenCounter): error=True, error_message=request_error, status_code=400, + rejected_before_provider_call=True, ) deployment = deployment or {} diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index 3c9c30b4a29..62503bf2aa6 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final, Literal from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import assert_never +from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for from litellm.types.llms.anthropic import AnthropicMessagesTextParam from litellm.types.llms.openai import AllMessageValues from litellm.types.llms.vertex_ai import ContentType @@ -70,27 +71,6 @@ def _is_function_tool(tool: Mapping[str, JsonValue]) -> bool: return isinstance(tool.get("name"), str) or "input_schema" in tool or "function" in tool -def _is_responses_function_tool(tool: Mapping[str, JsonValue]) -> bool: - return tool.get("type") == "function" and "function" not in tool and "input_schema" not in tool - - -def _chat_tool(tool: Mapping[str, JsonValue]) -> tuple[Mapping[str, object], ...]: - from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter - from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, - ) - - if _is_responses_function_tool(tool): - chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - [dict(tool)] # pyright: ignore[reportArgumentType] # validated JSON dict in the Responses function tool shape - ) - return tuple(chat_tools) - anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai( - [dict(tool)] # pyright: ignore[reportArgumentType] # validated JSON dict in an Anthropic or chat tool shape - ) - return tuple(anthropic_tools) - - def _map_function_tools( tools: object | None, ) -> tuple[Mapping[str, object], ...] | None | VertexGeminiCountContextError: @@ -107,7 +87,7 @@ def _map_function_tools( from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig chat_tools: Final = tuple( - chat_tool for tool in function_tools for chat_tool in _chat_tool(tool) + chat_tool for tool in function_tools for chat_tool in chat_completion_tools_for(tool) ) # comprehension-ok: flattens each tool's translation try: mapped_tools: Final = VertexGeminiConfig()._map_function( # pyright: ignore[reportPrivateUsage] # the tool mapping inference uses @@ -158,6 +138,7 @@ def count_context_error_response( error=True, error_message=error.error_message, status_code=400, + rejected_before_provider_call=True, ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d89c8baf9b0..9635cf3ce72 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19337,7 +19337,13 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "token_count_framing": 97 + }, + "databricks/system.ai.deepseek-v4-flash-0731": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 97 }, "databricks/databricks-deepseek-v4-pro-0813": { "cache_creation_input_token_cost": 1.31999e-06, @@ -20328,7 +20334,13 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 100 + }, + "databricks/system.ai.kimi-k3": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 100 }, "databricks/databricks-llama-4-maverick": { "cache_creation_input_token_cost": 5.0001e-07, @@ -75642,7 +75654,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "global.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 3.75e-06, @@ -75662,7 +75676,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "us.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 4.125e-06, @@ -75682,7 +75698,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "openrouter/xiaomi/mimo-v2.6-flash": { "cache_read_input_token_cost": 2.8e-09, diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 210bd5a8bfd..98e1cf92026 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -41,6 +41,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) +from litellm.proxy.common_utils.token_count_settings import count_source_fields router: Final = APIRouter() @@ -366,7 +367,7 @@ async def count_tokens( log_llm_api_exception(e, litellm_call_id) return _anthropic_error_json_response(proxy_exception_from_route_error(e), request) else: - return {"input_tokens": token_response.total_tokens} + return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)} @router.post( diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py new file mode 100644 index 00000000000..6f805616820 --- /dev/null +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -0,0 +1,164 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass, replace +from types import MappingProxyType +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, ValidationError + +import litellm +from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.proxy._types import ProxyException +from litellm.types.token_count import TokenCountSourceFields, TokenCountTokenizer +from litellm.types.utils import TokenCountResponse + +CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404}) +TOKENIZER_LOAD_WAIT_SECONDS: Final = 5.0 +TOKENIZER_LOAD_REUSE_SECONDS: Final = 60.0 +TOKENIZER_LOAD_MAX_RUNNING: Final = 4 + + +class _UpstreamErrorDetail(BaseModel): + model_config = ConfigDict(frozen=True, extra="allow") + message: str + + +class _UpstreamErrorEnvelope(BaseModel): + model_config = ConfigDict(frozen=True, extra="allow") + type: Literal["error"] + error: _UpstreamErrorDetail + + +def _joined(message: str, hint: str) -> str: + return f"{message.rstrip().rstrip('.')}. {hint}" + + +def _with_hint(message: str, hint: str) -> str: + try: + envelope: Final = _UpstreamErrorEnvelope.model_validate_json(message) + except ValidationError: + return _joined(AnthropicExceptionMapping.extract_error_message(message), hint) + detail: Final = envelope.error.model_copy(update={"message": _joined(envelope.error.message, hint)}) + return envelope.model_copy(update={"error": detail}).model_dump_json() + + +@dataclass(frozen=True, slots=True) +class TokenCountModel: + model: str + region: str | None + + +@dataclass(frozen=True, slots=True) +class TokenCountSettings: + framing: int | None + count_model: TokenCountModel | None + tokenizer: TokenCountTokenizer | None = None + + +@dataclass(frozen=True, slots=True) +class ProviderCountFailure: + status_code: int + message: str + rejected_before_provider_call: bool = False + + @property + def cannot_count_model(self) -> bool: + return not self.rejected_before_provider_call and self.status_code in CANNOT_COUNT_MODEL_STATUSES + + def with_hint(self, hint: str) -> "ProviderCountFailure": + return replace(self, message=_with_hint(self.message, hint)) + + def to_proxy_exception(self) -> ProxyException: + return ProxyException( + message=self.message, + type="token_counting_error", + param="model", + code=self.status_code, + ) + + +def cost_map_entry(model: str, custom_llm_provider: str | None = None) -> Mapping[str, object]: + try: + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models + return MappingProxyType({}) + + +def _framing(value: object) -> int | None: + match value: + case bool(): + return None + case int() if value >= 0: + return value + case float() if value.is_integer() and value >= 0: + return int(value) + case _: + return None + + +def _non_blank(value: object) -> str | None: + return value if isinstance(value, str) and value.strip() else None + + +def _count_model(value: Mapping[str, object]) -> TokenCountModel | None: + model: Final = _non_blank(value.get("model")) + return None if model is None else TokenCountModel(model=model, region=_non_blank(value.get("region"))) + + +def _geography(region: str) -> str: + return region.rsplit("-", 2)[0] + + +def _within_deployment_geography( + count_model: TokenCountModel | None, deployment_region: Callable[[], str] +) -> TokenCountModel | None: + if count_model is None or count_model.region is None or count_model.model.startswith("global."): + return count_model + if _geography(deployment_region()) == _geography(count_model.region): + return count_model + return replace(count_model, region=None) + + +def _resolved_count_model( + configured: object, default: object, deployment_region: Callable[[], str] +) -> TokenCountModel | None: + configured_mapping: Final = as_str_mapping(configured) + if configured_mapping is not None: + return _count_model(configured_mapping) + default_mapping: Final = as_str_mapping(default) + return ( + None + if default_mapping is None + else _within_deployment_geography(_count_model(default_mapping), deployment_region) + ) + + +def _tokenizer(value: object) -> TokenCountTokenizer | None: + match value: + case "kimi_k3": + return "kimi_k3" + case _: + return None + + +def resolve_token_count_settings( + configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str] +) -> TokenCountSettings: + framing: Final = _framing(configured.get("token_count_framing")) + return TokenCountSettings( + framing=framing if framing is not None else _framing(default.get("token_count_framing")), + count_model=_resolved_count_model( + configured.get("token_count_model"), default.get("token_count_model"), deployment_region + ), + tokenizer=_tokenizer(default.get("token_count_tokenizer")), + ) + + +def count_source_fields(response: TokenCountResponse) -> TokenCountSourceFields: + match response.count_source: + case None: + return {} + case "token_count_model": + return {"litellm_count_source": "token_count_model", "litellm_count_model": response.model_used} + case source: + return {"litellm_count_source": source} diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 818fd1ad516..5fbe780379e 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -18,6 +18,7 @@ from litellm.proxy.common_request_processing import ( resolve_litellm_call_id, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.types.llms.vertex_ai import TokenCountDetailsResponse @@ -306,10 +307,12 @@ async def google_count_tokens(request: Request, model_name: str): ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500), ) original_response: Final[dict] = token_response.original_response or {} - return TokenCountDetailsResponse( - totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0, - promptTokensDetails=original_response.get("promptTokensDetails", []), - ) + counted: Final[TokenCountDetailsResponse] = { + "totalTokens": original_response.get("totalTokens") or token_response.total_tokens or 0, + "promptTokensDetails": original_response.get("promptTokensDetails", []), + **count_source_fields(token_response), + } + return counted # ============================================================ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ee1e8647bf2..8369e79596c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -26,6 +26,7 @@ from collections.abc import ( MutableMapping, Sequence, ) +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from itertools import chain from types import MappingProxyType, UnionType @@ -47,6 +48,7 @@ from typing import ( ) import anyio +import anyio.lowlevel import websockets import websockets.exceptions from pydantic import BaseModel, Json, JsonValue, TypeAdapter, ValidationError @@ -87,6 +89,13 @@ from litellm.litellm_core_utils.litellm_logging import ( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + HUGGING_FACE_SOURCE, + KimiK3Encodings, + kimi_k3_anthropic_tokens, + kimi_k3_template_tokens, + kimi_k3_tokens, +) from litellm.proxy._types import ( UI_TEAM_ID, CallbackDelete, @@ -133,6 +142,15 @@ from litellm.proxy.common_utils.callback_utils import ( strip_callback_config, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body +from litellm.proxy.common_utils.token_count_settings import ( + TOKENIZER_LOAD_MAX_RUNNING, + TOKENIZER_LOAD_REUSE_SECONDS, + TOKENIZER_LOAD_WAIT_SECONDS, + ProviderCountFailure, + TokenCountSettings, + cost_map_entry, + resolve_token_count_settings, +) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.router_utils.access_windows import access_windows_config_error from litellm.router_utils.add_retry_fallback_headers import ( @@ -278,6 +296,7 @@ from litellm.constants import ( CLI_SSO_SESSION_TTL_SECONDS, DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, + DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MODEL_CREATED_AT_TIME, GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_ADMIN_NAME, @@ -308,8 +327,10 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, + as_str_mapping, drop_params_flag, get_litellm_metadata_from_kwargs, + iter_items, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -880,7 +901,13 @@ from litellm.types.secret_managers.main import ( KeyManagementSettings, KeyManagementSystem, ) -from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer, RawRequestTypedDict, StandardLoggingPayload +from litellm.types.utils import ( + CredentialItem, + CustomHuggingfaceTokenizer, + RawRequestTypedDict, + SelectTokenizerResponse, + StandardLoggingPayload, +) from litellm.types.utils import ModelInfo as ModelMapInfo from litellm.utils import _add_custom_logger_callback_to_specific_event @@ -13614,24 +13641,35 @@ async def _try_provider_token_count( provider_counter: "BaseTokenCounter", custom_llm_provider: str | None, model_to_use: str, + region: str | None, messages: list | None, contents: list | None, - deployment: dict[str, Any] | None, + deployment: Mapping[str, object] | None, request_model: str, tools: list | None = None, system: str | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, -) -> Optional["TokenCountResponse"]: - """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.""" +) -> Union["TokenCountResponse", ProviderCountFailure, None]: + """Attempt provider-specific token counting. Returns None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): return None + deployment_in_region: Final = ( + None + if deployment is None + else dict(deployment) + if region is None + else { + **deployment, + "litellm_params": {**(as_str_mapping(deployment.get("litellm_params")) or {}), "aws_region_name": region}, + } + ) try: result: Final = await provider_counter.count_tokens( model_to_use=model_to_use or "", messages=messages, contents=contents, - deployment=deployment, + deployment=deployment_in_region, request_model=request_model, tools=tools, system=system, @@ -13639,24 +13677,325 @@ async def _try_provider_token_count( output_format=output_format, ) except httpx.HTTPStatusError as e: - error_message: Final = getattr(e, "message", None) or str(e) - status_code: Final = getattr(e, "status_code", None) or e.response.status_code - raise ProxyException( - message=error_message, - type="token_counting_error", - param="model", - code=status_code, + return ProviderCountFailure( + status_code=getattr(e, "status_code", None) or e.response.status_code, + message=getattr(e, "message", None) or str(e), ) if result is not None and result.error is True: - raise ProxyException( + return ProviderCountFailure( + status_code=result.status_code or 500, message=result.error_message or "Token counting failed", - type="token_counting_error", - param="model", - code=result.status_code or 500, + rejected_before_provider_call=result.rejected_before_provider_call, ) return result +async def _provider_token_count( + count_with: Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]], + model_to_use: str, + settings: TokenCountSettings, + local_estimate_allowed: bool, +) -> Union["TokenCountResponse", ProviderCountFailure, None]: + """Count with the provider, using the token_count_model in its region when one is set. + + Returns None when the count should be estimated locally, or the provider's failure when the framed local + estimate stands in for it and that failure is the answer if the estimate cannot read the request.""" + count_model: Final = settings.count_model + counted: Final = await ( + count_with(model_to_use, None) if count_model is None else count_with(count_model.model, count_model.region) + ) + if counted is None: + return None + if isinstance(counted, TokenCountResponse): + if count_model is None: + return counted.model_copy(update={"count_source": "provider"}) + return counted.model_copy(update={"count_source": "token_count_model", "model_used": count_model.model}) + failure: Final = ( + counted + if count_model is None or count_model.region is None + else counted.with_hint( + f"This model counts tokens with {count_model.model} in {count_model.region}. " + "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" + ) + ) + if count_model is None and failure.cannot_count_model and settings.framing is not None and local_estimate_allowed: + return failure + raise failure.to_proxy_exception() + + +@dataclass(frozen=True, slots=True) +class _LocalTokenEstimate: + total_tokens: int + tokenizer_type: str + + +def _is_fork_guard_error(error: BaseException) -> bool: + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking + + return isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)) + + +def _is_tokenizer_panic(error: BaseException) -> bool: + return type(error).__module__ == "pyo3_runtime" and type(error).__name__ == "PanicException" + + +def _answers_provider_error(error: BaseException) -> bool: + return (isinstance(error, Exception) or _is_tokenizer_panic(error)) and not _is_fork_guard_error(error) + + +def _provider_error_after_failed_estimate( + provider_failure: ProviderCountFailure, error: BaseException +) -> ProxyException: + verbose_proxy_logger.warning( + "Local token estimate failed after the provider's count error %s, answering that error: %s", + provider_failure.status_code, + type(error).__name__, + ) + return provider_failure.to_proxy_exception() + + +async def _estimate_or_provider_error(estimate: Awaitable[int], provider_failure: ProviderCountFailure) -> int: + try: + return await estimate + except BaseException as e: + if not _answers_provider_error(e): + raise + raise _provider_error_after_failed_estimate(provider_failure, e) from e + + +@dataclass(frozen=True, slots=True) +class _CustomTokenizerKey: + identifier: str + revision: str + auth_token: str | None + + +class _TokenizerLoadsBusy(Exception): + pass + + +@dataclass(frozen=True, slots=True) +class _TokenizerLoadOutcome: + tokenizer: SelectTokenizerResponse | BaseException + finished_at: float + + +@dataclass(frozen=True, slots=True) +class _SharedTokenizerLoad: + outcome: asyncio.Task[_TokenizerLoadOutcome] + started_at: float + + def reusable(self, now: float) -> bool: + return not self.outcome.done() or now - self.outcome.result().finished_at < TOKENIZER_LOAD_REUSE_SECONDS + + +async def _loaded_tokenizer( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter +) -> SelectTokenizerResponse | BaseException: + try: + return await asyncify(litellm.utils._select_tokenizer, limiter=threads)( + model=model, custom_tokenizer=custom_tokenizer + ) + except asyncio.CancelledError: + raise + except BaseException as e: # noqa: BLE001 # every count sharing this load decides what it answers, pyo3 panics included + return e + + +async def _tokenizer_load_outcome( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter +) -> _TokenizerLoadOutcome: + tokenizer: Final = await _loaded_tokenizer(model, custom_tokenizer, threads) + return _TokenizerLoadOutcome(tokenizer=tokenizer, finished_at=asyncio.get_running_loop().time()) + + +@dataclass(slots=True) +class _TokenizerLoads: + threads: anyio.CapacityLimiter + running: int = 0 + + def start( + self, model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, loop: asyncio.AbstractEventLoop + ) -> _SharedTokenizerLoad | None: + if self.running >= self.threads.total_tokens: + return None + self.running += 1 + outcome: Final = loop.create_task(_tokenizer_load_outcome(model, custom_tokenizer, self.threads)) + outcome.add_done_callback(self._finished) + return _SharedTokenizerLoad(outcome=outcome, started_at=loop.time()) + + def _finished(self, _outcome: asyncio.Task[_TokenizerLoadOutcome]) -> None: + self.running -= 1 + + +_TOKENIZER_LOADS: Final = anyio.lowlevel.RunVar[_TokenizerLoads]("litellm_proxy_tokenizer_loads") + + +def _tokenizer_loads_for_this_loop() -> _TokenizerLoads: + existing: Final = _TOKENIZER_LOADS.get(None) + if existing is not None: + return existing + created: Final = _TokenizerLoads(threads=anyio.CapacityLimiter(TOKENIZER_LOAD_MAX_RUNNING)) + _TOKENIZER_LOADS.set(created) + return created + + +@dataclass(slots=True) +class _TokenizerLoadSlot: + model: str + custom_tokenizer: CustomHuggingfaceTokenizer | None + current: _SharedTokenizerLoad | None = None + + def load(self, loads: _TokenizerLoads, loop: asyncio.AbstractEventLoop) -> _SharedTokenizerLoad | None: + current: Final = self.current + if current is not None and current.reusable(loop.time()): + return current + self.current = loads.start(self.model, self.custom_tokenizer, loop) + return self.current + + +@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) +def _tokenizer_load_slot( + model: str, custom_tokenizer: _CustomTokenizerKey | None, loop: asyncio.AbstractEventLoop +) -> _TokenizerLoadSlot: + return _TokenizerLoadSlot( + model=model, + custom_tokenizer=None + if custom_tokenizer is None + else CustomHuggingfaceTokenizer( + identifier=custom_tokenizer.identifier, + revision=custom_tokenizer.revision, + auth_token=custom_tokenizer.auth_token, + ), + ) + + +async def _tokenizer_after_provider_error( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None +) -> SelectTokenizerResponse | BaseException: + loop: Final = asyncio.get_running_loop() + try: + load: Final = _tokenizer_load_slot( + model, + None + if custom_tokenizer is None + else _CustomTokenizerKey( + identifier=custom_tokenizer["identifier"], + revision=custom_tokenizer["revision"], + auth_token=custom_tokenizer["auth_token"], + ), + loop, + ).load(_tokenizer_loads_for_this_loop(), loop) + if load is None: + return _TokenizerLoadsBusy() + outcome: Final = await asyncio.wait_for( + asyncio.shield(load.outcome), + timeout=max(load.started_at + TOKENIZER_LOAD_WAIT_SECONDS - loop.time(), 0), + ) + except Exception as e: # noqa: BLE001 # a tokenizer that cannot be read, loaded in time, or loaded at all answers the provider's error + return e + return outcome.tokenizer + + +async def _tokenizer_or_provider_error( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, provider_failure: ProviderCountFailure +) -> SelectTokenizerResponse: + tokenizer: Final = await _tokenizer_after_provider_error(model, custom_tokenizer) + if not isinstance(tokenizer, BaseException): + return tokenizer + if not _answers_provider_error(tokenizer): + raise tokenizer.with_traceback(None) + raise _provider_error_after_failed_estimate(provider_failure, tokenizer) from tokenizer + + +async def _local_token_estimate( + model: str, + prompt: str | None, + messages: Sequence[AllMessageValues] | None, + tools: list[ChatCompletionToolParam] | None, + custom_tokenizer: CustomHuggingfaceTokenizer | None, + provider_failure: ProviderCountFailure | None, +) -> _LocalTokenEstimate: + tokenizer: Final = ( + await asyncify(litellm.utils._select_tokenizer)(model=model, custom_tokenizer=custom_tokenizer) + if provider_failure is None + else await _tokenizer_or_provider_error(model, custom_tokenizer, provider_failure) + ) + count: Final = offload_token_count(litellm.token_counter)( + model=model, + text=prompt, + messages=messages, + tools=tools, + custom_tokenizer=tokenizer, + ) + total_tokens: Final = await ( + count if provider_failure is None else _estimate_or_provider_error(count, provider_failure) + ) + return _LocalTokenEstimate(total_tokens=total_tokens, tokenizer_type=str(tokenizer["type"])) + + +_KIMI_K3_ENCODINGS: Final = KimiK3Encodings( + source=HUGGING_FACE_SOURCE, retry_after_seconds=TOKENIZER_LOAD_REUSE_SECONDS +) + + +async def _kimi_k3_count( + request: TokenCountRequest, + messages_format: TokenCountMessagesFormat, + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam] | None, +) -> int | None: + encoding: Final = await _KIMI_K3_ENCODINGS.encoding(TOKENIZER_LOAD_WAIT_SECONDS) + if encoding is None: + return None + match messages_format: + case "anthropic": + return await offload_token_count(kimi_k3_anthropic_tokens)( + encoding, model, request.model_dump(include={"messages", "system", "tools"}, exclude_none=True) + ) + case "openai": + return await offload_token_count(kimi_k3_tokens)(encoding, model, messages, tools) + + +class _TokenCountDeployment(TypedDict, total=False): + litellm_params: ReadOnly[Mapping[str, object]] + model_info: ReadOnly[Mapping[str, object]] + + +_TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment) + + +def _model_group_deployment(router: Router | None, model_group: str) -> Mapping[str, object] | None: + group: Final = () if router is None else router._get_all_deployments(model_name=model_group) + return group[0] if group else None + + +def _non_empty_str(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _deployment_aws_region(deployment: _TokenCountDeployment) -> str: + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + params: Final = deployment.get("litellm_params", {}) + return BaseAWSLLM()._get_aws_region_name( + {"aws_region_name": params.get("aws_region_name")}, model=_non_empty_str(params.get("model")) + ) + + +def _token_count_settings_model(deployment: _TokenCountDeployment, request_model: str) -> tuple[str, str | None]: + model_info: Final = deployment.get("model_info", {}) + params: Final = deployment.get("litellm_params", {}) + base_model: Final = _non_empty_str(model_info.get("base_model")) or _non_empty_str(params.get("base_model")) + if base_model is not None: + return base_model, None + model: Final = _non_empty_str(params.get("model")) + if model is None: + return request_model, None + return model, _non_empty_str(params.get("custom_llm_provider")) + + _PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( {litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value} ) @@ -13667,6 +14006,14 @@ def _drops_params(deployment: Mapping[str, object] | None) -> bool: return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) +_ESTIMABLE_GOOGLE_PART_KEYS: Final = frozenset({"text", "functionCall", "functionResponse", "thoughtSignature"}) + + +def _estimable_contents(contents: Sequence[Mapping[str, object]]) -> bool: + parts: Final = chain.from_iterable(iter_items(content.get("parts") or ()) for content in contents) + return all(set(as_str_mapping(part) or ()) <= _ESTIMABLE_GOOGLE_PART_KEYS for part in parts) + + def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -13735,6 +14082,14 @@ async def count_request_tokens( model_to_use: str = ( litellm_model_name if isinstance(litellm_model_name, str) and litellm_model_name else request.model ) + deployment_view: Final = _TOKEN_COUNT_DEPLOYMENT.validate_python( + deployment or _model_group_deployment(llm_router, request.model) or {} + ) + 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)), + deployment_region=partial(_deployment_aws_region, deployment_view), + ) # Try provider-specific token counting first - only for non-direct requests (from provider endpoints) provider_counter: BaseTokenCounter | None = None @@ -13754,23 +14109,32 @@ async def count_request_tokens( code=400, ) prompt_only: Final = bool(prompt) and conversation_empty - if provider_counter is not None: - prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES - result: Final = await _try_provider_token_count( - provider_counter=provider_counter, - custom_llm_provider=custom_llm_provider, + 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 + ladder: Final = ( + None + if provider_counter is None + else await _provider_token_count( + count_with=partial( + _try_provider_token_count, + provider_counter, + custom_llm_provider, + messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages, + contents=None if prompt_as_message else contents, + deployment=deployment, + request_model=request.model, + tools=tools, + system=system, + messages_format=messages_format, + output_format=output_format, + ), model_to_use=model_to_use, - messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages, - contents=None if prompt_as_message else contents, - deployment=deployment, - request_model=request.model, - tools=tools, - system=system, - messages_format=messages_format, - output_format=output_format, + settings=settings, + local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable, ) - if result is not None: - return result + ) + if isinstance(ladder, TokenCountResponse): + return ladder # Check if token counter is disabled before fallback if litellm.disable_token_counter is True: @@ -13788,11 +14152,6 @@ async def count_request_tokens( CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)( - model=model_to_use, custom_tokenizer=custom_tokenizer - ) - - tokenizer_used: Final = str(_tokenizer_used["type"]) system_message: Final = _system_message(system) prompt_turn: Final = 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 @@ -13805,18 +14164,38 @@ async def count_request_tokens( counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None ) - total_tokens: Final = await offload_token_count(litellm.token_counter)( + framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0 + kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None + exact: Final = ( + await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, counted_tools) + if kimi_k3_messages is not None and custom_tokenizer is None + else None + ) + if exact is not None: + return TokenCountResponse( + total_tokens=exact + framing, + request_model=request.model, + model_used=model_to_use, + tokenizer_type="huggingface_tokenizer", + count_source="local_estimate", + ) + estimate: Final = await _local_token_estimate( model=model_to_use, - text=prompt if counted_messages is None else None, + prompt=prompt if counted_messages is None else None, messages=counted_messages, tools=counted_tools, - custom_tokenizer=_tokenizer_used, + custom_tokenizer=custom_tokenizer, + provider_failure=ladder, + ) + template: Final = ( + 0 if kimi_k3_messages is None else kimi_k3_template_tokens(len(kimi_k3_messages), bool(counted_tools)) ) return TokenCountResponse( - total_tokens=total_tokens, + total_tokens=estimate.total_tokens + template + framing, request_model=request.model, model_used=model_to_use, - tokenizer_type=tokenizer_used, + tokenizer_type=estimate.tokenizer_type, + count_source="local_estimate", ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 021c741b2bc..d33d8a83bfa 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -42,6 +42,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, openai_error_type_for_status, ) +from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -49,6 +50,7 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ) from litellm.types.responses.main import DeleteResponseResult +from litellm.types.token_count import TokenCountSourceFields from litellm.types.utils import TokenCountResponse if TYPE_CHECKING: @@ -1074,7 +1076,7 @@ class _ResponsesApiErrorBody(TypedDict): error: ReadOnly[_ResponsesApiErrorDetail] -class _ResponsesInputTokensResult(TypedDict): +class _ResponsesInputTokensResult(TokenCountSourceFields): object: ReadOnly[str] input_tokens: ReadOnly[int] @@ -1289,6 +1291,7 @@ async def responses_input_tokens( result: Final[_ResponsesInputTokensResult] = { "object": "response.input_tokens", "input_tokens": token_response.total_tokens, + **count_source_fields(token_response), } return result diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index ce51e46ef15..28fdf3acf67 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -7,6 +7,7 @@ from typing_extensions import ( ) from litellm.types.llms.openai import EmbeddingInput +from litellm.types.token_count import TokenCountSourceFields # Gemini supports nested-list inputs (e.g. [["text", "image"]]) as an explicit # opt-in for combined embeddings — a provider-specific extension of the @@ -308,12 +309,13 @@ class UsageMetadata(TypedDict, total=False): candidatesTokensDetails: list[PromptTokensDetails] # Alternative key name used in some responses -class TokenCountDetailsResponse(TypedDict): +class TokenCountDetailsResponse(TokenCountSourceFields): """ Response structure for token count details with modality breakdown. Example: - {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]} + {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}], + 'litellm_count_source': 'provider'} """ totalTokens: int diff --git a/litellm/types/token_count.py b/litellm/types/token_count.py new file mode 100644 index 00000000000..3b21401066c --- /dev/null +++ b/litellm/types/token_count.py @@ -0,0 +1,16 @@ +from typing import Literal, TypeAlias + +from typing_extensions import NotRequired, ReadOnly, TypedDict + +TokenCountSource: TypeAlias = Literal["provider", "token_count_model", "local_estimate"] +TokenCountTokenizer: TypeAlias = Literal["kimi_k3"] + + +class TokenCountModelInfo(TypedDict): + model: ReadOnly[str] + region: ReadOnly[NotRequired[str]] + + +class TokenCountSourceFields(TypedDict, total=False): + litellm_count_source: ReadOnly[TokenCountSource] + litellm_count_model: ReadOnly[str] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d3f856697e3..2cce951ca60 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -55,6 +55,7 @@ from litellm.types.llms.base import ( LiteLLMPydanticObjectBase, ) from litellm.types.mcp import MCPServerCostInfo +from litellm.types.token_count import TokenCountModelInfo, TokenCountSource, TokenCountTokenizer from ..litellm_core_utils.core_helpers import map_finish_reason, process_response_headers from . import litellm_params as _litellm_params @@ -211,6 +212,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False): vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"] | None] bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None bedrock_converse_supports_strict_tools: bool | None + token_count_framing: ReadOnly[int | None] + token_count_model: ReadOnly[TokenCountModelInfo | None] + token_count_tokenizer: ReadOnly[TokenCountTokenizer | None] class SearchContextCostPerQuery(TypedDict, total=False): @@ -3826,10 +3830,13 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"}) +DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset({"token_count_framing", "token_count_model"}) + SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__) - frozenset(CustomPricingLiteLLMParams.model_fields) - DEPLOYMENT_SCOPED_PRICING_FIELDS + - DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS ) @@ -4226,6 +4233,7 @@ class TokenCountResponse(LiteLLMPydanticObjectBase): request_model: str model_used: str tokenizer_type: str + count_source: TokenCountSource | None = None original_response: dict | None = None """ Original Response from upstream API call - if an API call was made for token counting @@ -4236,6 +4244,7 @@ class TokenCountResponse(LiteLLMPydanticObjectBase): HTTP status code from the token counting API (e.g., 200 for success, 429 for rate limit, 400 for bad request) """ status_code: int | None = None + rejected_before_provider_call: bool = Field(default=False, exclude=True) class CustomHuggingfaceTokenizer(TypedDict): diff --git a/litellm/utils.py b/litellm/utils.py index e91460b5987..64ecfcac98c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6327,6 +6327,9 @@ def _get_model_info_helper( uses_embed_content=_model_info.get("uses_embed_content", None), supports_image_size=_model_info.get("supports_image_size", None), supported_audio_formats=_model_info.get("supported_audio_formats", None), + token_count_framing=_model_info.get("token_count_framing", None), + token_count_model=_model_info.get("token_count_model", None), + token_count_tokenizer=_model_info.get("token_count_tokenizer", None), vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None), ) for cost_key, cost_value in _model_info.items(): diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d89c8baf9b0..9635cf3ce72 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19337,7 +19337,13 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "token_count_framing": 97 + }, + "databricks/system.ai.deepseek-v4-flash-0731": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 97 }, "databricks/databricks-deepseek-v4-pro-0813": { "cache_creation_input_token_cost": 1.31999e-06, @@ -20328,7 +20334,13 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 100 + }, + "databricks/system.ai.kimi-k3": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 100 }, "databricks/databricks-llama-4-maverick": { "cache_creation_input_token_cost": 5.0001e-07, @@ -75642,7 +75654,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "global.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 3.75e-06, @@ -75662,7 +75676,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "us.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 4.125e-06, @@ -75682,7 +75698,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "openrouter/xiaomi/mimo-v2.6-flash": { "cache_read_input_token_cost": 2.8e-09, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 35624045fdf..b99f8ee6f98 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -1083,6 +1083,36 @@ "additionalProperties": false } }, + "token_count_framing": { + "type": "integer", + "minimum": 0, + "description": "Tokens added once per request to a local count estimate, for the provider's chat framing." + }, + "token_count_model": { + "type": "object", + "required": [ + "model" + ], + "additionalProperties": false, + "properties": { + "model": { + "type": "string", + "pattern": "\\S" + }, + "region": { + "type": "string", + "pattern": "\\S" + } + }, + "description": "Model that counts tokens for this one, and the AWS region its count API runs in." + }, + "token_count_tokenizer": { + "type": "string", + "description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.", + "enum": [ + "kimi_k3" + ] + }, "tpm": { "type": "integer", "minimum": 0, 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 a4b5c15ad54..b34b0b6b526 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 @@ -280,7 +280,7 @@ def test_messages_count_tokens_lifts_the_leading_system_run( ) -> None: with _listening(deployment) as peer: response: Final = _count(gateway, count_request(deployment.model, case)) - assert _payload(response) == {"input_tokens": _COUNT}, response.text + assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text assert _bodies(peer, deployment.provider) == (expected_count_body(_MODEL, case),) @@ -335,7 +335,11 @@ def test_responses_input_tokens_lifts_instructions(deployment: _Deployment, gate "/v1/responses/input_tokens", {"model": deployment.model, "input": USER_TEXT, "instructions": INSTRUCTION}, ) - assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _COUNT}, response.text + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _COUNT, + "litellm_count_source": "provider", + }, response.text assert _bodies(peer, deployment.provider) == (_STRING_BODY,) @@ -348,7 +352,11 @@ def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_ite "/v1/responses/input_tokens", {"model": deployment.model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION}, ) - assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _COUNT}, response.text + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _COUNT, + "litellm_count_source": "provider", + }, response.text assert _bodies(peer, deployment.provider) == ( {"model": _MODEL, "messages": [USER], "system": [LIFTED, {"type": "text", "text": REMINDER}]}, ) @@ -379,7 +387,7 @@ def test_messages_count_tokens_forwards_tools_beside_the_lifted_system( ) -> None: with _listening(deployment) as peer: response: Final = _count(gateway, {**count_request(deployment.model, STRING_CASE), "tools": TOOLS}) - assert _payload(response) == {"input_tokens": _COUNT}, response.text + assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text assert _bodies(peer, deployment.provider) == ({**_STRING_BODY, "tools": TOOLS},) @@ -392,7 +400,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place( "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], } response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": _COUNT}, response.text + assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text assert _bodies(peer, deployment.provider) == ( {"model": _MODEL, "messages": [USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], "system": [LIFTED]}, ) @@ -424,7 +432,7 @@ def test_messages_count_tokens_repeated_request_lifts_each_time(deployment: _Dep answers: Final = tuple( _payload(_count(gateway, count_request(deployment.model, STRING_CASE))) for _ in range(2) ) - assert answers == ({"input_tokens": _COUNT},) * 2 + assert answers == ({"input_tokens": _COUNT, "litellm_count_source": "provider"},) * 2 assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * 2 @@ -438,7 +446,7 @@ def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_ 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 + assert _payload(follow_up) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, follow_up.text assert _bodies(peer, deployment.provider) == (_STRING_BODY,) @@ -453,7 +461,7 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value( content=f'{{"model": "{deployment.model}", "messages": {first}, "messages": {last}}}', headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, ) - assert _payload(response) == {"input_tokens": _COUNT}, response.text + assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text assert _bodies(peer, deployment.provider) == (_STRING_BODY,) 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 bc70c5d3b21..e20a397a708 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -385,7 +385,7 @@ def test_messages_count_tokens_counts_through_mantle_when_the_runtime_cannot( ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) @@ -401,7 +401,7 @@ def test_messages_count_tokens_without_system_or_tools_sends_a_bare_body_to_mant ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_BARE,) @@ -436,7 +436,11 @@ def test_gemini_count_tokens_route_reports_the_mantle_total(counting_proxy: Owne ): model: Final = _deployment(scenario, runtime.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) - assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text + assert _payload(response) == { + "totalTokens": _MANTLE_COUNT, + "promptTokensDetails": [], + "litellm_count_source": "provider", + }, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == ({"model": _OPUS_BASE, "messages": _GEMINI_MESSAGES},) @@ -452,7 +456,11 @@ def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock ): model: Final = _deployment(scenario, runtime.url, _SONNET) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) - assert _payload(response) == {"totalTokens": _RUNTIME_COUNT, "promptTokensDetails": []}, response.text + assert _payload(response) == { + "totalTokens": _RUNTIME_COUNT, + "promptTokensDetails": [], + "litellm_count_source": "provider", + }, response.text assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) assert mantle.drain() == () @@ -467,7 +475,7 @@ def test_runtime_404_on_claude_falls_back_to_mantle(counting_proxy: OwnedProxy, ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) @@ -527,7 +535,11 @@ def test_google_count_tokens_forwards_system_and_tools_to_mantle( ): model: Final = _deployment(scenario, runtime.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) - assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text + assert _payload(response) == { + "totalTokens": _MANTLE_COUNT, + "promptTokensDetails": [], + "litellm_count_source": "provider", + }, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == ( _mantle_body( @@ -572,7 +584,11 @@ def test_google_count_tokens_normalizes_uppercase_parameters_for_mantle( ): model: Final = _deployment(scenario, runtime.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) - assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text + assert _payload(response) == { + "totalTokens": _MANTLE_COUNT, + "promptTokensDetails": [], + "litellm_count_source": "provider", + }, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == ( _mantle_body( @@ -602,7 +618,11 @@ def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy response: Final = gateway.request( "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"} ) - assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _MANTLE_COUNT, + "litellm_count_source": "provider", + }, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_BARE,) @@ -647,7 +667,10 @@ def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProx gateway.scenario() as scenario, ): model: Final = _deployment(scenario, runtime.url, _SONNET) - assert _payload(_count(gateway, _full(model))) == {"input_tokens": _RUNTIME_COUNT} + assert _payload(_count(gateway, _full(model))) == { + "input_tokens": _RUNTIME_COUNT, + "litellm_count_source": "provider", + } detailed: Final = gateway.request( "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} ) @@ -735,7 +758,11 @@ def test_responses_input_tokens_with_instructions_counts_through_mantle( "/v1/responses/input_tokens", {"model": model, "input": "Count this message", "instructions": "Be terse"}, ) - assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _MANTLE_COUNT, + "litellm_count_source": "provider", + }, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(system=[{"type": "text", "text": "Be terse"}]),) @@ -752,7 +779,7 @@ def test_messages_count_tokens_lifts_the_leading_system_run_for_mantle( ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, count_request(model, case)) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (expected_count_body(_OPUS_BASE, case),) @@ -865,7 +892,11 @@ def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_ite "/v1/responses/input_tokens", {"model": model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION}, ) - assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _MANTLE_COUNT, + "litellm_count_source": "provider", + }, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED, {"type": "text", "text": REMINDER}]),) @@ -881,7 +912,7 @@ def test_messages_count_tokens_forwards_tools_beside_the_lifted_system_to_mantle ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, {**count_request(model, STRING_CASE), "tools": _TOOLS}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED], tools=_TOOLS),) @@ -901,7 +932,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place_for_mant "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], } response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == ( _mantle_body(messages=[USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], system=[LIFTED]), @@ -954,7 +985,7 @@ def test_messages_count_tokens_repeated_request_lifts_the_leading_system_each_ti ): model: Final = _deployment(scenario, runtime.url) answers: Final = tuple(_payload(_count(gateway, count_request(model, STRING_CASE))) for _ in range(2)) - assert answers == ({"input_tokens": _MANTLE_COUNT},) * 2 + assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2 assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * 2 @@ -975,7 +1006,9 @@ def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_ 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 + assert _payload(follow_up) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, ( + follow_up.text + ) assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) @@ -997,7 +1030,7 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value_for_m content=f'{{"model": "{model}", "messages": {first}, "messages": {last}}}', headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, ) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) @@ -1022,7 +1055,9 @@ def test_disabled_token_counter_counts_a_leading_system_through_mantle( ) with wire_server(_mantle(_strict), port=mantle_port) as strict: counted: Final = _count(owned.gateway, count_request(_OWNED_OPUS, STRING_CASE)) - assert _payload(counted) == {"input_tokens": _MANTLE_COUNT}, counted.text + assert _payload(counted) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, ( + counted.text + ) assert _mantle_bodies(strict) == (_MANTLE_LIFTED,) assert len(_runtime_count_targets(runtime)) == 1 @@ -1222,7 +1257,7 @@ def test_messages_count_tokens_forwards_a_5kb_system_verbatim_to_mantle( ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(system=system),) @@ -1243,7 +1278,7 @@ def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}', headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, ) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 (sent,) = mantle.drain() assert _mantle_requests((sent,)) == (_MANTLE_FULL,) @@ -1261,7 +1296,7 @@ def test_messages_count_tokens_leaves_empty_tools_and_system_out_of_the_mantle_b ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "tools": [], "system": ""}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_BARE,) @@ -1277,7 +1312,7 @@ def test_messages_count_tokens_repeated_request_reaches_both_peers_each_time( ): model: Final = _deployment(scenario, runtime.url) answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2)) - assert answers == ({"input_tokens": _MANTLE_COUNT},) * 2 + assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2 assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 assert _mantle_bodies(mantle) == (_MANTLE_BARE,) * 2 @@ -1306,9 +1341,15 @@ def test_disabled_token_counter_surfaces_the_mantle_error_instead_of_counting_lo assert _REJECTION in refused.text and "input_tokens" not in refused.text, refused.text assert _mantle_bodies(refusing) == (_MANTLE_FULL,) with wire_server(_mantle(), port=mantle_port) as counting: - assert _payload(_count(owned.gateway, _full(_OWNED_OPUS))) == {"input_tokens": _MANTLE_COUNT} + assert _payload(_count(owned.gateway, _full(_OWNED_OPUS))) == { + "input_tokens": _MANTLE_COUNT, + "litellm_count_source": "provider", + } assert _mantle_bodies(counting) == (_MANTLE_FULL,) - assert _payload(_count(owned.gateway, _full(_OWNED_SONNET))) == {"input_tokens": _RUNTIME_COUNT} + assert _payload(_count(owned.gateway, _full(_OWNED_SONNET))) == { + "input_tokens": _RUNTIME_COUNT, + "litellm_count_source": "provider", + } rejected: Final = _count(owned.gateway, _full(_OWNED_NOVA)) assert rejected.status_code == 400, rejected.text assert _UNSUPPORTED in rejected.text and "input_tokens" not in rejected.text, rejected.text diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index 1c4b1d1a6e5..d5c841d9335 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -399,7 +399,7 @@ def test_messages_count_tokens_forwards_system_and_tools(gateway: Gateway, syste with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system, "tools": _TOOLS},) @@ -407,7 +407,7 @@ def test_messages_count_tokens_without_system_or_tools_sends_bare_body(gateway: with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text assert _counted_bodies(wire) == (_PEER_BARE,) @@ -446,7 +446,11 @@ def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway response: Final = gateway.request( "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"} ) - assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _PEER_COUNT, + "litellm_count_source": "provider", + }, response.text assert _counted_bodies(wire) == (_PEER_BARE,) @@ -501,7 +505,7 @@ def test_llama_count_still_reaches_the_meta_count_route(gateway: Gateway) -> Non with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{_LLAMA_BACKEND}") response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text (sent,) = _count_bodies(wire.drain(), _LLAMA_COUNT_TARGET) assert sent["messages"] == _MESSAGES, sent @@ -619,7 +623,7 @@ def test_messages_count_tokens_forwards_an_empty_tools_list(gateway: Gateway) -> with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, {**_bare(model), "tools": []}) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": []},) @@ -627,7 +631,7 @@ def test_messages_count_tokens_forwards_an_empty_system_string(gateway: Gateway) with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, {**_bare(model), "system": ""}) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text assert _counted_bodies(wire) == ({**_PEER_BARE, "system": ""},) @@ -635,7 +639,7 @@ def test_messages_count_tokens_leaves_null_system_and_tools_out(gateway: Gateway with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, {**_bare(model), "system": None, "tools": None}) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text assert _counted_bodies(wire) == (_PEER_BARE,) @@ -653,7 +657,7 @@ def test_messages_count_tokens_forwards_a_5kb_system_verbatim(gateway: Gateway) with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, {**_bare(model), "system": system}) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system},) @@ -666,7 +670,7 @@ def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}', headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, ) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text (sent,) = _count_requests(wire.drain()) assert _count_bodies((sent,)) == (_PEER_FULL,) assert (sent.body.count(b'"system"'), sent.body.count(b'"tools"')) == (1, 1), sent.body @@ -699,7 +703,7 @@ def test_count_tokens_location_override_targets_the_count_region(gateway: Gatewa gateway, scenario, wire.url, vertex_location="global", vertex_count_tokens_location="europe-west1" ) response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text target: Final = _COUNT_TARGET.replace(f"/locations/{_LOCATION}/", "/locations/europe-west1/") assert _count_bodies(wire.drain(), target) == (_PEER_BARE,) @@ -708,7 +712,7 @@ def test_messages_count_tokens_repeated_request_reaches_the_peer_each_time(gatew with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2)) - assert answers == ({"input_tokens": _PEER_COUNT},) * 2 + assert answers == ({"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"},) * 2 assert _counted_bodies(wire) == (_PEER_BARE,) * 2 @@ -725,7 +729,7 @@ def test_disabled_token_counter_surfaces_provider_failures_instead_of_counting_l ) with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: counted: Final = _count(owned.gateway, _full(_OWNED_MODEL)) - assert _payload(counted) == {"input_tokens": _PEER_COUNT}, counted.text + assert _payload(counted) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, counted.text rejected: Final = _count( owned.gateway, {**_full(_OWNED_MODEL), "messages": [{"role": "user", "content": _REJECT_TEXT}]} ) diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index 1cc03838f8d..b5145ea49d4 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -399,9 +399,11 @@ def test_in_flight_count_tokens_does_not_reserve_key_budget_away_from_a_completi [{"role": "user", "content": prompt}] ] assert count.status_code == 200, count.text - assert count.json() == {"totalTokens": 12, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}]}, ( - count.text - ) + assert count.json() == { + "totalTokens": 12, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}], + "litellm_count_source": "provider", + }, count.text provider_calls: Final = wire.drain() assert [(call.method, call.target) for call in provider_calls] == [ ("POST", "/v1beta/models/gemini-3.8-flash:countTokens") diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py index 114bbb184c6..2777d18a399 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -4,10 +4,11 @@ Test to verify the Google GenAI proxy API endpoints """ from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest +from litellm.types.utils import TokenCountResponse def _build_test_client(): @@ -88,9 +89,7 @@ def test_google_stream_generate_content_endpoint(): assert init_kwargs["data"]["_litellm_raw_sse_stream"] is True assert init_kwargs["data"]["_litellm_skip_openai_stream_done"] is True assert init_kwargs["data"]["model"] == "test-model" - assert init_kwargs["data"]["contents"] == [ - {"role": "user", "parts": [{"text": "Hello"}]} - ] + assert init_kwargs["data"]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] def test_google_generate_content_data_flows_through_processor(): @@ -123,9 +122,7 @@ def test_google_generate_content_data_flows_through_processor(): data = mock_init.call_args.kwargs["data"] assert data["model"] == "test-model" assert data["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] - assert data["systemInstruction"] == { - "parts": [{"text": "Your name is Doodle."}] - } + assert data["systemInstruction"] == {"parts": [{"text": "Your name is Doodle."}]} # generationConfig arrives intact here; the rename to `config` is # done downstream in route_request (see test_route_llm_request). assert data["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] @@ -160,12 +157,17 @@ def test_google_count_tokens_unchanged(): except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - fake_response = MagicMock() - fake_response.original_response = { - "totalTokens": 7, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}], - } - fake_response.total_tokens = 7 + fake_response = TokenCountResponse( + total_tokens=7, + request_model="test-model", + model_used="test-model", + tokenizer_type="gemini_api", + original_response={ + "totalTokens": 7, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}], + }, + count_source="provider", + ) with patch( "litellm.proxy.proxy_server.token_counter", @@ -180,6 +182,7 @@ def test_google_count_tokens_unchanged(): assert response.status_code == 200 body = response.json() assert body["totalTokens"] == 7 + assert body["litellm_count_source"] == "provider" def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope(): @@ -212,9 +215,14 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - fake_response: Final = MagicMock() - fake_response.original_response = {"input_tokens": 2167} - fake_response.total_tokens = 2167 + fake_response: Final = TokenCountResponse( + total_tokens=2167, + request_model="claude-opus-4-8", + model_used="us.anthropic.claude-opus-5-5", + tokenizer_type="bedrock_api", + original_response={"input_tokens": 2167}, + count_source="token_count_model", + ) with patch( "litellm.proxy.proxy_server.token_counter", @@ -230,3 +238,5 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs() body: Final = response.json() assert body["totalTokens"] == 2167 assert body["promptTokensDetails"] == [] + assert body["litellm_count_source"] == "token_count_model" + assert body["litellm_count_model"] == "us.anthropic.claude-opus-5-5" diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 5f7807650e1..b7f6865ab7e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -4239,6 +4239,54 @@ class TestModelInfoCostMapEchoFilter: assert "litellm_provider" not in info assert "supported_openai_params" not in info + @pytest.fixture + def bedrock_sonnet_5_5_mapped_as_on_main(self) -> Iterator[None]: + import litellm + + entry: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"}, + } + with pytest.MonkeyPatch.context() as patched: + patched.setattr(litellm, "model_cost", {**litellm.model_cost, "us.anthropic.claude-sonnet-5-5": entry}) + litellm.utils._invalidate_model_cost_lowercase_map() + yield + litellm.utils._invalidate_model_cost_lowercase_map() + + @pytest.mark.parametrize( + ("model", "field", "edited"), + [ + ("bedrock/us.anthropic.claude-sonnet-5-5", "token_count_model", {}), + ("bedrock/us.moonshotai.kimi-k3", "token_count_framing", 0), + ], + ) + def test_token_count_setting_is_stored_only_when_the_operator_changed_it( + self, model: str, field: str, edited: object, bedrock_sonnet_5_5_mapped_as_on_main: None + ) -> None: + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + entry: Final = litellm.get_model_info(model) + assert entry[field] not in (None, edited), entry + db_model: Final = Deployment( + model_name="count-settings", + litellm_params=LiteLLM_Params(model=model), + model_info=ModelInfo(id="dep-echo-count"), + ) + + echoed: Final = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count"})), + ) + changed: Final = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count", field: edited})), + ) + + assert field not in json.loads(echoed["model_info"]) + assert json.loads(changed["model_info"])[field] == edited + def test_metadata_without_a_cost_map_key_is_persisted(self): import litellm diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py index fc5ab9475eb..ee165d7b204 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py @@ -52,6 +52,7 @@ def test_token_counter_happy_path(client, auth_as, patched_token_counter): "request_model": "gpt-4", "model_used": "gpt-4", "tokenizer_type": "openai_tokenizer", + "count_source": "local_estimate", "original_response": None, "error": False, "error_message": None, diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index 86ad6c8d9d9..4913d2b0d1d 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -54,6 +54,7 @@ async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload( "request_model": "azure_ai/claude-sonnet-4-5", "model_used": "claude-sonnet-4-5", "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, "original_response": payload, "error": False, "error_message": None, @@ -98,6 +99,7 @@ async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_v "request_model": "azure_ai/claude-sonnet-4-5", "model_used": "claude-sonnet-4-5", "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, "original_response": None, "error": True, "error_message": str(rejection.value), @@ -126,6 +128,7 @@ async def test_count_payload_that_is_not_an_object_becomes_an_error_response_nam "request_model": "azure_ai/claude-sonnet-4-5", "model_used": "claude-sonnet-4-5", "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, "original_response": None, "error": True, "error_message": str(failed_lookup.value), @@ -156,6 +159,7 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_ "request_model": "azure_ai/claude-sonnet-4-5", "model_used": "claude-sonnet-4-5", "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, "original_response": None, "error": True, "error_message": upstream_text, @@ -180,6 +184,7 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_ "request_model": "azure_ai/claude-sonnet-4-5", "model_used": "claude-sonnet-4-5", "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, "original_response": None, "error": True, "error_message": f"CountTokens processing error: {undecodable.value}", diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py new file mode 100644 index 00000000000..ffddb09cca7 --- /dev/null +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py @@ -0,0 +1,524 @@ +import asyncio +import base64 +import hashlib +import json +import logging +import re +import threading +from collections.abc import Iterable +from itertools import product +from pathlib import Path +from typing import Final + +import httpx +import pytest +import respx +import tiktoken + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + KimiK3Encodings, + KimiK3TokenizerSource, + kimi_k3_anthropic_tokens, + kimi_k3_tokens, + load_kimi_k3_encoding, +) +from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig + +MODEL: Final = "us.moonshotai.kimi-k3" +_CONTROL: Final = re.compile(r"<\|(?:open|close|sep|end_of_msg)\|>") +_WEATHER: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather in °C", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}, "days": {"type": "integer"}}, + "required": ["city"], + }, + }, +} + + +def _call(call_id: str, arguments: dict[str, object]) -> dict[str, object]: + return {"id": call_id, "type": "function", "function": {"name": "get_weather", "arguments": json.dumps(arguments)}} + + +def _written_ranks(path: Path, tokens: Iterable[bytes]) -> KimiK3TokenizerSource: + path.write_bytes(b"\n".join(base64.b64encode(token) + b" %d" % rank for rank, token in enumerate(tokens))) + return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + + +def _single_bytes() -> tuple[bytes, ...]: + return tuple(bytes((value,)) for value in range(256)) + + +@pytest.fixture(scope="module") +def byte_source(tmp_path_factory: pytest.TempPathFactory) -> KimiK3TokenizerSource: + return _written_ranks(tmp_path_factory.mktemp("bytes") / "tiktoken.model", _single_bytes()) + + +@pytest.fixture(scope="module") +def byte_encoding(byte_source: KimiK3TokenizerSource) -> tiktoken.Encoding: + loaded: Final = load_kimi_k3_encoding(byte_source) + assert loaded is not None + return loaded + + +@pytest.fixture(scope="module") +def pair_encoding(tmp_path_factory: pytest.TempPathFactory) -> tiktoken.Encoding: + pairs: Final = (bytes(pair) for pair in product(range(256), repeat=2)) + source: Final = _written_ranks(tmp_path_factory.mktemp("pairs") / "tiktoken.model", (*_single_bytes(), *pairs)) + loaded: Final = load_kimi_k3_encoding(source) + assert loaded is not None + return loaded + + +def _template_count(encoding: tiktoken.Encoding, rendered: str) -> int: + texts: Final = _CONTROL.split(rendered) + return len(_CONTROL.findall(rendered)) + sum(len(encoding.encode_ordinary(text)) for text in texts) + + +_GENERATION_PROMPT: Final = '<|open|>message role="assistant"<|sep|><|open|>think<|sep|>' + + +# Expected renders follow the Kimi K3 chat template in moonshotai/Kimi-K3 at revision f831ab6, which matched +# Bedrock's prompt_tokens for us.moonshotai.kimi-k3 in 19 live calls on 2026-10-16 +@pytest.mark.parametrize( + ("messages", "tools", "rendered"), + [ + pytest.param( + [{"role": "user", "content": "Hello there"}], + None, + '<|open|>message role="user"<|sep|>Hello there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT, + id="user", + ), + pytest.param( + [ + { + "role": "system", + "content": [{"type": "text", "text": "Be terse."}, {"type": "text", "text": "Answer in French."}], + }, + {"role": "user", "content": [{"type": "text", "text": "Hi"}, {"type": "text", "text": " there"}]}, + ], + None, + '<|open|>message role="system"<|sep|>Be terse.Answer in French.<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="user"<|sep|>Hi there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT, + id="text_parts_join", + ), + pytest.param( + [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "And 3+3?"}, + ], + None, + '<|open|>message role="user"<|sep|>What is 2+2?<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>' + "<|open|>response<|sep|>4<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|>" + '<|open|>message role="user"<|sep|>And 3+3?<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT, + id="assistant_turn", + ), + pytest.param( + [ + {"role": "user", "content": "Weather in Paris and Oslo?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [_call("c1", {"city": "Paris", "days": 3}), _call("c2", {"city": "Oslo"})], + }, + {"role": "tool", "tool_call_id": "c2", "content": "Snow"}, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ], + [_WEATHER], + '<|open|>message role="system" type="tool-declare"<|sep|># Tools\n' + "Here are the available tools, described in JSONSchema.\n\n```json\n" + '[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":' + '{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},' + '"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="user"<|sep|>Weather in Paris and Oslo?<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>' + "<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>" + '<|open|>call tool="get_weather" index="1"<|sep|>' + '<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|>' + '<|open|>argument key="days" type="number"<|sep|>3<|close|>argument<|sep|><|close|>call<|sep|>' + '<|open|>call tool="get_weather" index="2"<|sep|>' + '<|open|>argument key="city" type="string"<|sep|>Oslo<|close|>argument<|sep|><|close|>call<|sep|>' + "<|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>" + '<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="tool" tool="get_weather" index="2"<|sep|>Snow<|close|>message<|sep|><|end_of_msg|>' + + _GENERATION_PROMPT, + id="tool_results_out_of_call_order", + ), + pytest.param( + [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + _call( + "c1", {"a&b": 'x"y', "flag": True, "none": None, "list": [1, "zé"], "object": {"k": True}} + ) + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "Rain"}, + ], + [_WEATHER], + '<|open|>message role="system" type="tool-declare"<|sep|># Tools\n' + "Here are the available tools, described in JSONSchema.\n\n```json\n" + '[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":' + '{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},' + '"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="user"<|sep|>Weather?<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>' + "<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>" + '<|open|>call tool="get_weather" index="1"<|sep|>' + '<|open|>argument key="a&b" type="string"<|sep|>x"y<|close|>argument<|sep|>' + '<|open|>argument key="flag" type="boolean"<|sep|>true<|close|>argument<|sep|>' + '<|open|>argument key="none" type="null"<|sep|>null<|close|>argument<|sep|>' + '<|open|>argument key="list" type="array"<|sep|>[1, "zé"]<|close|>argument<|sep|>' + '<|open|>argument key="object" type="object"<|sep|>{"k": true}<|close|>argument<|sep|>' + "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>" + '<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Rain<|close|>message<|sep|><|end_of_msg|>' + + _GENERATION_PROMPT, + id="argument_types_and_escaping", + ), + ], +) +@pytest.mark.parametrize("encoding_fixture", ["byte_encoding", "pair_encoding"]) +def test_count_is_the_kimi_k3_chat_template_render_of_the_bedrock_request( + request: pytest.FixtureRequest, + encoding_fixture: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + rendered: str, +): + encoding: Final[tiktoken.Encoding] = request.getfixturevalue(encoding_fixture) + + assert kimi_k3_tokens(encoding, MODEL, messages, tools) == _template_count(encoding, rendered) + + +def _parallel_calls(count: int) -> list[dict[str, object]]: + ids: Final = tuple(f"c{number}" for number in range(1, count + 1)) + return [ + {"role": "user", "content": "Weather?"}, + {"role": "assistant", "content": None, "tool_calls": [_call(call_id, {"city": "Paris"}) for call_id in ids]}, + *({"role": "tool", "tool_call_id": call_id, "content": "Sun"} for call_id in ids), + ] + + +def test_tenth_parallel_call_and_its_result_are_numbered_ten(byte_encoding: tiktoken.Encoding): + nine: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(9), [_WEATHER]) + ten: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(10), [_WEATHER]) + assert nine is not None and ten is not None + + assert ten - nine == _template_count( + byte_encoding, + '<|open|>call tool="get_weather" index="10"<|sep|>' + '<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|><|close|>call<|sep|>' + '<|open|>message role="tool" tool="get_weather" index="10"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>', + ) + + +@pytest.mark.parametrize( + "text", + [ + pytest.param("The quick brown fox, jumps over 12345 lazy dogs. " * 60, id="prose"), + pytest.param("a" + "1234567890" * 250 + "1", id="long_digit_run"), + pytest.param("--- " * 600, id="symbols"), + pytest.param("def f(x):\n return x # it's 42\n\n" * 80, id="code"), + pytest.param("word next\t\ttab\n \n" * 120, id="whitespace"), + pytest.param("naïve café, “quotes” & 中文汉字，日本語のテキスト 😀🎉\r\n" * 40, id="unicode"), + pytest.param("We're sure they'd say it's DON'T or I'M " * 60, id="contractions"), + pytest.param("abc123def4567ghi89 " * 120, id="letters_and_digits"), + ], +) +def test_long_text_counts_what_the_tokenizer_counts_for_the_whole_text(pair_encoding: tiktoken.Encoding, text: str): + def counted(content: str) -> int | None: + return kimi_k3_tokens(pair_encoding, MODEL, [{"role": "user", "content": content}], None) + + short: Final = counted("x") + assert short is not None + + assert counted(text) == short - len(pair_encoding.encode_ordinary("x")) + len(pair_encoding.encode_ordinary(text)) + + +@pytest.mark.parametrize( + ("messages", "tools"), + [ + pytest.param( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, + ], + } + ], + None, + id="image", + ), + pytest.param( + [{"role": "user", "content": "Weather?"}, {"role": "tool", "tool_call_id": "c1", "content": "Sun"}], + [_WEATHER], + id="tool_result_without_its_call", + ), + ], +) +def test_request_the_template_cannot_render_is_not_counted( + pair_encoding: tiktoken.Encoding, messages: list[dict[str, object]], tools: list[dict[str, object]] | None +): + assert kimi_k3_tokens(pair_encoding, MODEL, messages, tools) is None + + +@pytest.mark.parametrize( + "tool", + [ + pytest.param( + { + "type": "function", + "name": "get_weather", + "description": "Get the weather in °C", + "parameters": _WEATHER["function"]["parameters"], + }, + id="responses_api", + ), + pytest.param( + { + "name": "get_weather", + "description": "Get the weather in °C", + "input_schema": _WEATHER["function"]["parameters"], + }, + id="anthropic_messages", + ), + ], +) +def test_tool_in_another_api_shape_counts_like_the_same_chat_completions_tool( + pair_encoding: tiktoken.Encoding, tool: dict[str, object] +): + messages: Final = [{"role": "user", "content": "Weather in Paris?"}] + + counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [tool]) + + assert counted is not None + assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER]) + + +@pytest.mark.parametrize( + "tool", + [ + pytest.param({"type": "web_search_20250305", "name": "web_search"}, id="anthropic_web_search"), + pytest.param({"type": "web_search"}, id="responses_web_search"), + pytest.param( + { + "type": "mcp", + "server_label": "docs", + "server_url": "https://mcp.example/mcp", + "require_approval": "never", + }, + id="mcp", + ), + ], +) +def test_tool_bedrock_inference_drops_counts_like_the_function_tools_alone( + pair_encoding: tiktoken.Encoding, tool: dict[str, object] +): + messages: Final = [{"role": "user", "content": "Weather in Paris?"}] + + counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, tool]) + + assert counted is not None + assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER]) + + +def test_responses_custom_tool_counts_like_the_function_tool_inference_sends(pair_encoding: tiktoken.Encoding): + messages: Final = [{"role": "user", "content": "Weather in Paris?"}] + custom_tool: Final = {"type": "custom", "name": "grammar", "description": "Free-form SQL"} + sent_to_inference, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + [dict(custom_tool)] + ) + + without_it: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER]) + + counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, custom_tool]) + + assert counted is not None and without_it is not None + assert counted > without_it + assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, *sent_to_inference]) + + +def test_counting_a_request_with_an_image_url_never_downloads_the_image( + pair_encoding: tiktoken.Encoding, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "user_url_validation", False) + image: Final = respx_mock.get("https://images.example/cat.png").mock(return_value=httpx.Response(200)) + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + {"type": "image_url", "image_url": {"url": "https://images.example/cat.png"}}, + ], + } + ] + + assert kimi_k3_tokens(pair_encoding, MODEL, messages, None) is None + assert not image.called + + +def test_anthropic_request_counts_like_the_same_openai_request(pair_encoding: tiktoken.Encoding): + anthropic: Final = kimi_k3_anthropic_tokens( + pair_encoding, + MODEL, + { + "system": "Be terse.", + "messages": [ + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "c1", "name": "get_weather", "input": {"city": "Paris"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "Sun"}]}, + ], + "tools": [ + { + "name": "get_weather", + "description": "Get the weather in °C", + "input_schema": _WEATHER["function"]["parameters"], + } + ], + }, + ) + openai: Final = kimi_k3_tokens( + pair_encoding, + MODEL, + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": None, "tool_calls": [_call("c1", {"city": "Paris"})]}, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ], + [_WEATHER], + ) + + assert openai is not None + assert anthropic == openai + + +@pytest.mark.parametrize("failure", ["checksum", "missing_file"]) +def test_tokenizer_that_fails_its_checksum_or_cannot_be_read_is_not_loaded( + byte_source: KimiK3TokenizerSource, tmp_path: Path, caplog: pytest.LogCaptureFixture, failure: str +): + source: Final = ( + KimiK3TokenizerSource(ranks_path=byte_source.ranks_path, sha256=hashlib.sha256(b"other").hexdigest()) + if failure == "checksum" + else KimiK3TokenizerSource(ranks_path=lambda: str(tmp_path / "missing"), sha256=byte_source.sha256) + ) + verbose_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + loaded: Final = load_kimi_k3_encoding(source) + finally: + verbose_logger.removeHandler(caplog.handler) + + assert loaded is None + assert any("Kimi K3 tokenizer unavailable" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_slow_tokenizer_load_is_not_waited_on_past_the_wait_and_is_shared_once_done( + byte_source: KimiK3TokenizerSource, +): + released: Final = threading.Event() + + def slow_ranks_path() -> str: + released.wait(10) + return byte_source.ranks_path() + + encodings: Final = KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=slow_ranks_path, sha256=byte_source.sha256), + retry_after_seconds=60, + clock=lambda: 0.0, + ) + + assert await encodings.encoding(wait_seconds=0) is None + released.set() + loaded: Final = await encodings.encoding(wait_seconds=5) + assert loaded is not None + assert await encodings.encoding(wait_seconds=0) is loaded + + +class _Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +async def _loaded_eventually(encodings: KimiK3Encodings) -> tiktoken.Encoding | None: + for _ in range(500): + if (loaded := await encodings.encoding(wait_seconds=0)) is not None: + return loaded + await asyncio.sleep(0.01) + return None + + +@pytest.mark.asyncio +async def test_failed_tokenizer_load_is_retried_in_the_background_only_after_the_retry_window( + byte_source: KimiK3TokenizerSource, tmp_path: Path +): + clock: Final = _Clock() + reads: Final[list[float]] = [] + released: Final = threading.Event() + + def ranks_path() -> str: + reads.append(clock.now) + if len(reads) == 1: + return str(tmp_path / "missing") + released.wait(10) + return byte_source.ranks_path() + + encodings: Final = KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=ranks_path, sha256=byte_source.sha256), + retry_after_seconds=60, + clock=clock, + ) + + assert await encodings.encoding(wait_seconds=5) is None + clock.now = 59.0 + assert await encodings.encoding(wait_seconds=5) is None + clock.now = 60.0 + assert await asyncio.wait_for(encodings.encoding(wait_seconds=5), timeout=1) is None + released.set() + assert await _loaded_eventually(encodings) is not None + assert reads == [0.0, 60.0] + + +@pytest.mark.asyncio +async def test_tokenizer_load_still_running_after_the_wait_is_logged( + byte_source: KimiK3TokenizerSource, caplog: pytest.LogCaptureFixture +): + released: Final = threading.Event() + + def hung_ranks_path() -> str: + released.wait(10) + return byte_source.ranks_path() + + encodings: Final = KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=hung_ranks_path, sha256=byte_source.sha256), retry_after_seconds=60 + ) + verbose_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + loaded: Final = await encodings.encoding(wait_seconds=0.05) + finally: + verbose_logger.removeHandler(caplog.handler) + released.set() + + assert loaded is None + assert any("Kimi K3 tokenizer still loading" in record.getMessage() for record in caplog.records) diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py index 65a9f6f044d..a45773b93fc 100644 --- a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -54,6 +54,7 @@ async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload( "request_model": "openai/gpt-4o", "model_used": "gpt-4o", "tokenizer_type": "openai_api", + "count_source": None, "original_response": payload, "error": False, "error_message": None, @@ -98,6 +99,7 @@ async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_v "request_model": "openai/gpt-4o", "model_used": "gpt-4o", "tokenizer_type": "openai_api", + "count_source": None, "original_response": None, "error": True, "error_message": str(rejection.value), @@ -126,6 +128,7 @@ async def test_count_payload_that_is_not_an_object_becomes_an_error_response_nam "request_model": "openai/gpt-4o", "model_used": "gpt-4o", "tokenizer_type": "openai_api", + "count_source": None, "original_response": None, "error": True, "error_message": str(failed_lookup.value), @@ -156,6 +159,7 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_ "request_model": "openai/gpt-4o", "model_used": "gpt-4o", "tokenizer_type": "openai_api", + "count_source": None, "original_response": None, "error": True, "error_message": upstream_text, @@ -382,6 +386,7 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_ "request_model": "openai/gpt-4o", "model_used": "gpt-4o", "tokenizer_type": "openai_api", + "count_source": None, "original_response": None, "error": True, "error_message": f"CountTokens processing error: {undecodable.value}", diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py new file mode 100644 index 00000000000..84a32a02a5c --- /dev/null +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -0,0 +1,190 @@ +import json +from typing import Final + +import pytest + +from litellm.proxy.common_utils.token_count_settings import ( + ProviderCountFailure, + TokenCountModel, + TokenCountSettings, + count_source_fields, + resolve_token_count_settings, +) +from litellm.types.utils import TokenCountResponse + +_OPUS_US_EAST_1: Final = TokenCountModel(model="anthropic.claude-opus-5-5", region="us-east-1") +_DEFAULT: Final = { + "token_count_framing": 97, + "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}, +} + + +def _in_us_west_2() -> str: + return "us-west-2" + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ({}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ( + {"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}}, + TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")), + ), + ( + {"token_count_model": {"model": "custom-model"}}, + TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + ), + ( + {"token_count_model": {"model": "custom-model", "region": " "}}, + TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + ), + ({"token_count_framing": 0}, TokenCountSettings(framing=0, count_model=_OPUS_US_EAST_1)), + ( + {"token_count_framing": None, "token_count_model": None}, + TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1), + ), + ( + {"token_count_framing": True, "token_count_model": "custom-model"}, + TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1), + ), + ({"token_count_framing": 78.0}, TokenCountSettings(framing=78, count_model=_OPUS_US_EAST_1)), + ({"token_count_framing": 1.5}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ({"token_count_framing": -1}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ({"token_count_framing": "78"}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ({"token_count_model": {}}, TokenCountSettings(framing=97, count_model=None)), + ({"token_count_model": ""}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ( + {"token_count_model": {"model": " ", "region": "us-east-1"}}, + TokenCountSettings(framing=97, count_model=None), + ), + ], +) +def test_config_overrides_the_cost_map_default_key_by_key( + configured: dict[str, object], expected: TokenCountSettings +) -> None: + assert ( + resolve_token_count_settings(configured=configured, default=_DEFAULT, deployment_region=_in_us_west_2) + == expected + ) + + +@pytest.mark.parametrize( + ("deployment_region", "configured", "expected_region"), + [ + ("us-west-2", {}, "us-east-1"), + ("eu-central-1", {}, None), + ("ca-central-1", {}, None), + ("us-gov-west-1", {}, None), + ( + "eu-central-1", + {"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}}, + "us-east-1", + ), + ], +) +def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geography( + deployment_region: str, configured: dict[str, object], expected_region: str | None +) -> None: + settings: Final = resolve_token_count_settings( + configured=configured, default=_DEFAULT, deployment_region=lambda: deployment_region + ) + assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region) + + +def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography() -> None: + settings: Final = resolve_token_count_settings( + configured={}, + default={"token_count_model": {"model": "global.anthropic.claude-opus-5-5", "region": "us-east-1"}}, + deployment_region=lambda: "eu-west-1", + ) + assert settings.count_model == TokenCountModel(model="global.anthropic.claude-opus-5-5", region="us-east-1") + + +def test_shipped_australian_count_region_stays_with_an_australian_deployment() -> None: + settings: Final = resolve_token_count_settings( + configured={}, + default={"token_count_model": {"model": "au.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}}, + deployment_region=lambda: "ap-southeast-2", + ) + assert settings.count_model == TokenCountModel(model="au.anthropic.claude-opus-5-5", region="ap-southeast-4") + + +def test_no_setting_anywhere_means_the_model_counts_itself_without_framing() -> None: + assert resolve_token_count_settings( + configured={}, default={}, deployment_region=_in_us_west_2 + ) == TokenCountSettings(framing=None, count_model=None) + + +@pytest.mark.parametrize( + ("value", "expected"), + [("kimi_k3", "kimi_k3"), ("cl100k_base", None), (None, None), (True, None)], +) +def test_only_a_tokenizer_litellm_ships_is_taken_from_the_cost_map(value: object, expected: str | None) -> None: + settings: Final = resolve_token_count_settings( + configured={}, default={"token_count_tokenizer": value}, deployment_region=_in_us_west_2 + ) + assert settings.tokenizer == expected + + +def _counted(source: str | None, model_used: str = "us.anthropic.claude-opus-5-5") -> TokenCountResponse: + return TokenCountResponse( + total_tokens=7, + request_model="sonnet", + model_used=model_used, + tokenizer_type="bedrock_api", + count_source=source, + ) + + +@pytest.mark.parametrize( + ("source", "expected"), + [ + (None, {}), + ("provider", {"litellm_count_source": "provider"}), + ("local_estimate", {"litellm_count_source": "local_estimate"}), + ( + "token_count_model", + {"litellm_count_source": "token_count_model", "litellm_count_model": "us.anthropic.claude-opus-5-5"}, + ), + ], +) +def test_count_source_fields_name_the_counting_model_only_when_token_count_model_counted( + source: str | None, expected: dict[str, str] +) -> None: + assert count_source_fields(_counted(source)) == expected + + +_UPSTREAM_404: Final = { + "type": "error", + "request_id": "req_upstream", + "error": {"type": "not_found_error", "message": "The model 'anthropic.claude-opus-5-5' does not exist"}, +} + + +@pytest.mark.parametrize( + ("message", "expected"), + [ + ( + json.dumps(_UPSTREAM_404), + { + **_UPSTREAM_404, + "error": { + "type": "not_found_error", + "message": "The model 'anthropic.claude-opus-5-5' does not exist. Enable eu-west-1", + }, + }, + ), + (json.dumps({"message": "Too many tokens"}), "Too many tokens. Enable eu-west-1"), + ("scripted 404", "scripted 404. Enable eu-west-1"), + ( + "Credential should be scoped to a valid region. ", + "Credential should be scoped to a valid region. Enable eu-west-1", + ), + ], +) +def test_hint_lands_in_the_upstream_error_message_and_keeps_its_envelope(message: str, expected: object) -> None: + hinted: Final = ProviderCountFailure(status_code=404, message=message).with_hint("Enable eu-west-1") + + assert (json.loads(hinted.message) if isinstance(expected, dict) else hinted.message) == expected + assert hinted.status_code == 404 diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 869e7eb60d2..86556f1c7bd 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2,14 +2,32 @@ # 1. Generate a Key, and use it to make a call +import asyncio +import base64 +import hashlib +import itertools import json -from typing import Final +import logging +import threading +import time +import traceback +from collections.abc import Callable, Coroutine, Iterator, Mapping +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import Final, NoReturn from unittest.mock import AsyncMock, MagicMock, patch +import anyio import httpx import pytest +import tiktoken from dotenv import load_dotenv from pydantic import BaseModel, ConfigDict, JsonValue +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.processors import TemplateProcessing load_dotenv() @@ -21,10 +39,20 @@ from fastapi.responses import JSONResponse import litellm from litellm import Router +from litellm._logging import verbose_proxy_logger from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler +from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError 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 ( + KIMI_K3_RANKS_SHA256, + KimiK3Encodings, + KimiK3TokenizerSource, + kimi_k3_template_tokens, + kimi_k3_tokens, + load_kimi_k3_encoding, +) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter @@ -32,8 +60,59 @@ from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.anthropic_endpoints.endpoints import ( count_tokens as anthropic_count_tokens, ) -from litellm.proxy.proxy_server import token_counter -from litellm.types.utils import TokenCountResponse +from litellm.proxy.common_utils.token_count_settings import ( + ProviderCountFailure, + TokenCountModel, + TokenCountSettings, + cost_map_entry, +) +from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request +from litellm.proxy.proxy_server import ( + _estimate_or_provider_error, + _provider_token_count, + _try_provider_token_count, + count_request_tokens, + token_counter, +) +from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +from litellm.types.utils import SelectTokenizerResponse, TokenCountMessagesFormat, TokenCountResponse + + +def _hugging_face_offline() -> NoReturn: + raise OSError("unit tests never download the Kimi K3 tokenizer") + + +@pytest.fixture(autouse=True) +def _kimi_k3_tokenizer_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "_KIMI_K3_ENCODINGS", + KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=_hugging_face_offline, sha256=KIMI_K3_RANKS_SHA256), + retry_after_seconds=60, + ), + ) + + +@pytest.fixture +def bedrock_sonnet_5_5_mapped_as_on_main() -> Iterator[None]: + entries: Final = { + name: { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "token_count_model": {"model": count_model, "region": "us-east-1"}, + } + for name, count_model in ( + ("anthropic.claude-sonnet-5-5", "anthropic.claude-opus-5-5"), + ("us.anthropic.claude-sonnet-5-5", "us.anthropic.claude-opus-5-5"), + ) + } + with pytest.MonkeyPatch.context() as patched: + patched.setattr(litellm, "model_cost", {**litellm.model_cost, **entries}) + litellm.utils._invalidate_model_cost_lowercase_map() + yield + litellm.utils._invalidate_model_cost_lowercase_map() @pytest.mark.asyncio @@ -1179,6 +1258,1285 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): proxy_server.count_request_tokens = original_token_counter +class _ScriptedCounter(BaseTokenCounter): + def __init__(self, replies: dict[str, TokenCountResponse], burst: int = 1) -> None: + self._replies = replies + self._burst = burst + self._burst_arrived = asyncio.Event() + self.calls: list[tuple[str, object, object, object]] = [] + self.regions: list[object] = [] + + def should_use_token_counting_api(self, custom_llm_provider: str | None = None) -> bool: + return True + + async def count_tokens( + self, + model_to_use: str, + messages: object, + contents: object, + deployment: object = None, + request_model: str = "", + tools: object = None, + system: object = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, + context_management: object = None, + ) -> TokenCountResponse | None: + self.calls.append((model_to_use, messages, tools, system)) + self.regions.append(None if deployment is None else deployment["litellm_params"].get("aws_region_name")) + if len(self.calls) >= self._burst: + self._burst_arrived.set() + await self._burst_arrived.wait() + return self._replies[model_to_use] + + +def _provider_reply( + model: str, total_tokens: int = 0, status_code: int | None = None, rejected_before_provider_call: bool = False +) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=total_tokens, + request_model="sonnet-5-5", + model_used=model, + tokenizer_type="bedrock_api", + error=status_code is not None, + error_message=None if status_code is None else f"scripted {status_code} for {model}", + status_code=status_code, + rejected_before_provider_call=rejected_before_provider_call, + ) + + +_COUNT_REQUEST = TokenCountRequest( + model="sonnet-5-5", + messages=[{"role": "user", "content": "hello"}], + tools=[{"name": "get_quote", "input_schema": {"type": "object"}}], + system="be brief", +) + + +_SONNET_DEPLOYMENT: Final = MappingProxyType( + {"litellm_params": {"model": "bedrock/us.anthropic.claude-sonnet-5-5", "aws_region_name": "us-west-2"}} +) +_OPUS_IN_US_EAST_1: Final = TokenCountModel(model="us.anthropic.claude-opus-5-5", region="us-east-1") + + +async def _ladder( + counter: _ScriptedCounter, settings: TokenCountSettings, local_estimate_allowed: bool = True +) -> TokenCountResponse | None: + return await _provider_token_count( + count_with=partial( + _try_provider_token_count, + counter, + "bedrock", + messages=_COUNT_REQUEST.messages, + contents=_COUNT_REQUEST.contents, + deployment=_SONNET_DEPLOYMENT, + request_model=_COUNT_REQUEST.model, + tools=_COUNT_REQUEST.tools, + system=_COUNT_REQUEST.system, + ), + model_to_use="us.anthropic.claude-sonnet-5-5", + settings=settings, + local_estimate_allowed=local_estimate_allowed, + ) + + +@pytest.mark.asyncio +async def test_model_without_token_count_model_counts_itself_in_the_deployment_region(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=447)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=10, count_model=None)) + + assert result is not None + assert (result.total_tokens, result.count_source) == (447, "provider") + assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-sonnet-5-5"], ["us-west-2"]) + + +@pytest.mark.parametrize( + ("count_model", "region"), + [ + (_OPUS_IN_US_EAST_1, "us-east-1"), + (TokenCountModel(model="us.anthropic.claude-opus-5-5", region=None), "us-west-2"), + ], +) +@pytest.mark.asyncio +async def test_token_count_model_counts_the_same_request_in_its_own_region_without_trying_the_model( + count_model: TokenCountModel, region: str +): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=None, count_model=count_model)) + + assert result is not None + assert (result.total_tokens, result.count_source, result.model_used) == ( + 447, + "token_count_model", + "us.anthropic.claude-opus-5-5", + ) + assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-opus-5-5"], [region]) + assert counter.calls[0][1:] == (_COUNT_REQUEST.messages, _COUNT_REQUEST.tools, _COUNT_REQUEST.system) + assert _SONNET_DEPLOYMENT["litellm_params"]["aws_region_name"] == "us-west-2" + + +@pytest.mark.parametrize("status_code", [401, 403, 429, 500]) +@pytest.mark.asyncio +async def test_errors_other_than_cannot_count_raise_even_with_framing(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=status_code)} + ) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=10, count_model=None)) + + assert (raised.value.code, raised.value.message) == (str(status_code), f"scripted {status_code} for sonnet") + + +@pytest.mark.parametrize("status_code", [400, 403, 404, 500]) +@pytest.mark.asyncio +async def test_token_count_model_failure_names_its_region_so_the_user_can_enable_it(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=None, count_model=_OPUS_IN_US_EAST_1)) + + assert raised.value.code == str(status_code) + assert raised.value.message.startswith(f"scripted {status_code} for opus") + assert "us.anthropic.claude-opus-5-5 in us-east-1" in raised.value.message + assert len(counter.calls) == 1 + + +@pytest.mark.asyncio +async def test_token_count_model_without_a_region_answers_its_error_unchanged(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=404)}) + + with pytest.raises(ProxyException) as raised: + await _ladder( + counter, + TokenCountSettings( + framing=None, count_model=TokenCountModel(model="us.anthropic.claude-opus-5-5", region=None) + ), + ) + + assert (raised.value.code, raised.value.message) == ("404", "scripted 404 for opus") + + +@pytest.mark.asyncio +async def test_uncountable_model_without_framing_raises(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=404)}) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=None, count_model=None)) + + assert raised.value.code == "404" + + +@pytest.mark.asyncio +async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=None)) + + assert isinstance(result, ProviderCountFailure) + assert len(counter.calls) == 1 + + +@pytest.mark.parametrize("status_code", [400, 404]) +@pytest.mark.asyncio +async def test_token_count_model_error_stands_even_with_framing(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=78, count_model=_OPUS_IN_US_EAST_1)) + + assert raised.value.code == str(status_code) + assert "us.anthropic.claude-opus-5-5 in us-east-1" in raised.value.message + + +@pytest.mark.asyncio +async def test_uncountable_model_with_framing_raises_when_local_counting_is_disabled(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=10, count_model=None), local_estimate_allowed=False) + + assert raised.value.code == "400" + + +@pytest.mark.asyncio +async def test_request_the_counter_rejects_before_calling_the_provider_answers_400_without_framing(): + counter: Final = _ScriptedCounter( + { + "us.anthropic.claude-sonnet-5-5": _provider_reply( + "sonnet", status_code=400, rejected_before_provider_call=True + ) + } + ) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=100, count_model=None)) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for sonnet") + + +@pytest.mark.parametrize("model", ["openai/gpt-5.4-mini", "gemini/gemini-3.8-flash"]) +@pytest.mark.asyncio +async def test_framing_keeps_the_400_for_blocks_the_provider_counter_cannot_convert( + monkeypatch: pytest.MonkeyPatch, model: str +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "framed", + "litellm_params": {"model": model, "api_key": "sk-unused-framed"}, + "model_info": {"token_count_framing": 100}, + } + ] + ), + ) + search_result: Final = {"type": "search_result", "source": "https://example.com", "title": "Notes", "content": []} + + with pytest.raises(ProxyException) as raised: + await count_request_tokens( + request=TokenCountRequest( + model="framed", + messages=[{"role": "user", "content": [{"type": "text", "text": "Summarize."}, search_result]}], + ), + call_endpoint=True, + messages_format="anthropic", + ) + + assert raised.value.code == "400" + assert "search_result" in raised.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "unreadable_body", + [ + pytest.param( + { + "messages": [ + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "new"}, + }, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_01", "content": []}, + ], + }, + ] + }, + id="web_search_blocks", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "What is new?"}], + "tools": [{"name": "search", "input_schema": {"type": "object", "properties": {"query": True}}}], + }, + id="boolean_subschema_tool", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "What is new?"}], + "tools": [ + {"name": "search", "input_schema": {"type": "object", "properties": {"tags": {"type": "array"}}}} + ], + }, + id="array_schema_without_items", + ), + pytest.param( + {"messages": [], "contents": [{"role": "user", "parts": [{"text": "What is new?"}]}]}, + id="google_contents", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "Describe this"}], + "contents": [ + { + "role": "user", + "parts": [ + {"text": "Describe this"}, + {"fileData": {"mimeType": "image/png", "fileUri": "gs://b/c.png"}}, + ], + } + ], + }, + id="google_contents_with_media_the_translation_dropped", + ), + pytest.param( + { + "messages": [ + {"role": "user", "content": "What is new?"}, + {"role": "assistant", "content": "Weighing the release notes"}, + ], + "contents": [ + {"role": "user", "parts": [{"text": "What is new?"}]}, + {"role": "model", "parts": [{"text": "Weighing the release notes", "thought": True}]}, + ], + }, + id="google_contents_with_a_thought_the_translation_counts_as_text", + ), + ], +) +async def test_framed_estimate_that_cannot_read_the_body_answers_the_provider_error( + monkeypatch: pytest.MonkeyPatch, unreadable_body: dict[str, list[dict[str, object]]] +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + with pytest.raises(ProxyException) as raised: + await token_counter(request=TokenCountRequest(model="kimi", **unreadable_body), call_endpoint=True) + readable: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for kimi") + assert readable.count_source == "local_estimate" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "contents", + [ + pytest.param([{"role": "user", "parts": [{"text": "What is new in the release notes?"}]}], id="text"), + pytest.param( + [ + {"role": "user", "parts": [{"text": "Weather in Paris?"}]}, + { + "role": "model", + "parts": [ + {"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}, "thoughtSignature": "c2ln"} + ], + }, + { + "role": "user", + "parts": [{"functionResponse": {"name": "get_weather", "response": {"forecast": "Sun"}}}], + }, + ], + id="sdk_agent_loop", + ), + pytest.param( + [ + {"role": "user", "parts": [{"text": "Hi"}]}, + {"role": "model", "parts": [{"text": "Hello!", "thoughtSignature": "c2ln"}]}, + {"role": "user", "parts": [{"text": "What is new?"}]}, + ], + id="sdk_chat_with_a_thought_signature", + ), + pytest.param([], id="empty"), + ], +) +async def test_framed_estimate_counts_google_contents_like_the_same_conversation_sent_as_messages( + monkeypatch: pytest.MonkeyPatch, contents: list[dict[str, object]] +): + _framed_kimi(monkeypatch) + google_request: Final = _google_count_tokens_request("kimi", {"contents": contents}) + + from_google: Final = await token_counter(request=google_request, call_endpoint=True) + from_messages: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=google_request.messages), call_endpoint=True + ) + + assert from_google.count_source == "local_estimate" + assert from_google.total_tokens == from_messages.total_tokens + + +@pytest.mark.asyncio +async def test_framed_deployment_whose_custom_tokenizer_cannot_load_answers_the_provider_error( + monkeypatch: pytest.MonkeyPatch, +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "kimi", + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": {"custom_tokenizer": {"identifier": "moonshotai/kimi-k3", "revision": "main"}}, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for kimi") + + +def _framed_kimis_with_custom_tokenizers( + monkeypatch: pytest.MonkeyPatch, tokenizers: Mapping[str, tuple[str, str]], burst: int = 1 +) -> None: + counter: Final = _ScriptedCounter( + { + "us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400), + "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=12), + }, + burst=burst, + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": {"model": "bedrock/us.anthropic.claude-sonnet-5-5"}, + "model_info": {"token_count_model": {}}, + }, + *( + { + "model_name": name, + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": { + "custom_tokenizer": {"identifier": identifier, "revision": revision, "auth_token": None} + }, + } + for name, (identifier, revision) in tokenizers.items() + ), + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + +def _framed_kimi_with_custom_tokenizer(monkeypatch: pytest.MonkeyPatch, identifier: str) -> None: + _framed_kimis_with_custom_tokenizers(monkeypatch, {"kimi": (identifier, "main")}) + + +async def _count_status(model: str = "kimi") -> str: + try: + response: Final = await token_counter( + request=TokenCountRequest(model=model, messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + except ProxyException as e: + return e.code + return response.count_source or "" + + +class _SkippableClockLoop(asyncio.SelectorEventLoop): + skipped_seconds: float = 0.0 + + def time(self) -> float: + return super().time() + self.skipped_seconds + + def skip(self, seconds: float) -> None: + self.skipped_seconds += seconds + + +def _run_on_skippable_clock( + scenario: Callable[[_SkippableClockLoop], Coroutine[object, object, tuple[str, ...]]], +) -> tuple[str, ...]: + loop: Final = _SkippableClockLoop() + try: + return loop.run_until_complete(scenario(loop)) + finally: + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + +async def _until_finished(tasks: tuple[asyncio.Future[str], ...], count: int) -> None: + for finished in itertools.islice(asyncio.as_completed(tasks), count): + await finished + + +def _tokenizer_that_panics_while_encoding() -> SelectTokenizerResponse: + template: Final = Tokenizer(WordLevel({"[UNK]": 0, "": 1}, unk_token="[UNK]")) + template.pre_tokenizer = Whitespace() + template.post_processor = TemplateProcessing(single=" $A", special_tokens=[("", 1)]) + spec: Final = json.loads(template.to_str()) + return litellm.utils.create_tokenizer( + json.dumps({**spec, "post_processor": {**spec["post_processor"], "special_tokens": {}}}) + ) + + +def test_hanging_tokenizer_download_answers_the_provider_error_after_the_wait_with_one_download_while_it_runs( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/hanging-tokenizer") + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + downloads: Final[list[str]] = [] + + def hanging_download(model: str, custom_tokenizer: object = None) -> NoReturn: + downloads.append(model) + download_started.set() + hub_released.wait(10) + raise OSError("hub unreachable") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hanging_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e-9) + + async def burst_past_the_wait_then_one_more(loop: _SkippableClockLoop) -> tuple[str, ...]: + try: + burst: Final = asyncio.gather(*(_count_status() for _ in range(60))) + await asyncio.to_thread(download_started.wait, 10) + loop.skip(31) + answered: Final = await asyncio.wait_for(burst, 8) + later: Final = await asyncio.wait_for(_count_status(), 8) + finally: + hub_released.set() + return (*answered, later) + + assert set(_run_on_skippable_clock(burst_past_the_wait_then_one_more)) == {"400"} + assert len(downloads) == 1 + + +@pytest.mark.asyncio +async def test_framed_count_waits_for_a_tokenizer_download_that_finishes_within_the_wait( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/slow-tokenizer") + real_select_tokenizer: Final = litellm.utils._select_tokenizer + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + + def slow_download(model: str, custom_tokenizer: object = None): + download_started.set() + hub_released.wait(10) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", slow_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + count: Final = asyncio.create_task(_count_status()) + try: + await asyncio.to_thread(download_started.wait, 10) + finally: + hub_released.set() + + assert await asyncio.wait_for(count, 10) == "local_estimate" + + +@pytest.mark.asyncio +async def test_failed_tokenizer_download_is_reused_until_the_retry_window_passes(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/flaky-tokenizer") + real_select_tokenizer: Final = litellm.utils._select_tokenizer + downloads: Final[list[str]] = [] + + def hub_down_once(model: str, custom_tokenizer: object = None): + downloads.append(model) + if len(downloads) == 1: + raise OSError("hub unreachable") + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hub_down_once) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e9) + within_window: Final = [await _count_status() for _ in range(3)] + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e-9) + after_window: Final = await _count_status() + + assert within_window == ["400", "400", "400"] + assert after_window == "local_estimate" + assert len(downloads) == 2 + + +def test_failed_tokenizer_load_is_reused_for_the_window_after_it_fails(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/slow-failing-tokenizer") + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + downloads: Final[list[str]] = [] + + def slow_failing_download(model: str, custom_tokenizer: object = None) -> NoReturn: + downloads.append(model) + download_started.set() + hub_released.wait(10) + raise OSError("hub unreachable") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", slow_failing_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 1000) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 60) + + async def fail_after_longer_than_the_window_then_retry(loop: _SkippableClockLoop) -> tuple[str, ...]: + first: Final = asyncio.ensure_future(_count_status()) + try: + await asyncio.to_thread(download_started.wait, 10) + loop.skip(120) + finally: + hub_released.set() + return await asyncio.wait_for(first, 8), await asyncio.wait_for(_count_status(), 8) + + assert _run_on_skippable_clock(fail_after_longer_than_the_window_then_retry) == ("400", "400") + assert len(downloads) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "second_tokenizer", + [("lit-tests/tokenizer-b", "main"), ("lit-tests/tokenizer-a", "v2")], + ids=["identifier", "revision"], +) +async def test_framed_deployments_with_different_custom_tokenizers_each_load_their_own( + monkeypatch: pytest.MonkeyPatch, second_tokenizer: tuple[str, str] +): + _framed_kimis_with_custom_tokenizers( + monkeypatch, {"kimi-a": ("lit-tests/tokenizer-a", "main"), "kimi-b": second_tokenizer} + ) + real_select_tokenizer: Final = litellm.utils._select_tokenizer + loaded: Final[list[object]] = [] + + def load(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + loaded.append(custom_tokenizer) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", load) + + statuses: Final = [await _count_status(name) for name in ("kimi-a", "kimi-b")] + + assert statuses == ["local_estimate", "local_estimate"] + assert loaded == [ + {"identifier": "lit-tests/tokenizer-a", "revision": "main", "auth_token": None}, + {"identifier": second_tokenizer[0], "revision": second_tokenizer[1], "auth_token": None}, + ] + + +@pytest.mark.asyncio +async def test_tokenizer_loads_past_the_cap_answer_the_provider_error_without_stalling_other_counts( + monkeypatch: pytest.MonkeyPatch, +): + tokenizers: Final = {f"kimi-{n}": (f"lit-tests/unreachable-tokenizer-{n}", "main") for n in range(5)} + _framed_kimis_with_custom_tokenizers(monkeypatch, tokenizers, burst=len(tokenizers)) + real_select_tokenizer: Final = litellm.utils._select_tokenizer + hub_released: Final = threading.Event() + downloads: Final[list[object]] = [] + + def held_download(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + downloads.append(custom_tokenizer) + hub_released.wait(10) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", held_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 2) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + monkeypatch.setattr(anyio.to_thread.current_default_thread_limiter(), "total_tokens", 2) + framed: Final = tuple(asyncio.ensure_future(_count_status(name)) for name in tokenizers) + try: + await asyncio.wait_for(_until_finished(framed, 3), 8) + unrelated: Final = await asyncio.wait_for(_count_status("sonnet"), 5) + finally: + hub_released.set() + statuses: Final = await asyncio.wait_for(asyncio.gather(*framed), 8) + + assert sorted(statuses) == ["400"] * 3 + ["local_estimate"] * 2 + assert unrelated == "provider" + assert len(downloads) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("panics_while", ["loading", "counting"]) +async def test_framed_count_answers_the_provider_error_when_the_tokenizer_panics( + monkeypatch: pytest.MonkeyPatch, panics_while: str +): + _framed_kimi_with_custom_tokenizer(monkeypatch, f"lit-tests/panicking-tokenizer-{panics_while}") + panicking: Final = _tokenizer_that_panics_while_encoding() + + def load(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + if panics_while == "loading": + panicking["tokenizer"].encode("hello world") + return panicking + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", load) + + assert await _count_status() == "400" + + +def _framed_kimi(monkeypatch: pytest.MonkeyPatch) -> None: + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + +@pytest.mark.asyncio +async def test_local_count_loads_the_tokenizer_again_after_a_failed_load(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/local-route-tokenizer") + real_select_tokenizer: Final = litellm.utils._select_tokenizer + downloads: Final[list[str]] = [] + + def hub_down_once(model: str, custom_tokenizer: object = None): + downloads.append(model) + if len(downloads) == 1: + raise OSError("hub unreachable") + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hub_down_once) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e9) + request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]) + + with pytest.raises(OSError, match="hub unreachable"): + await token_counter(request=request) + retried: Final = await token_counter(request=request) + + assert retried.count_source == "local_estimate" and retried.total_tokens > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("reuse_seconds", "tokenizer_types"), + [(1e9, ("openai_tokenizer", "openai_tokenizer")), (1e-9, ("openai_tokenizer", "huggingface_tokenizer"))], + ids=["within_the_window", "after_the_window"], +) +async def test_framed_count_loads_the_tokenizer_again_only_after_the_window_of_a_load_that_fell_back( + monkeypatch: pytest.MonkeyPatch, reuse_seconds: float, tokenizer_types: tuple[str, str] +): + _framed_kimi(monkeypatch) + real_select_tokenizer: Final = litellm.utils._select_tokenizer + hub_tokenizer: Final = litellm.utils.create_tokenizer( + Tokenizer(WordLevel({"[UNK]": 0}, unk_token="[UNK]")).to_str() + ) + loads: Final[list[str]] = [] + + def hub_up_on_the_second_load(model: str, custom_tokenizer: object = None): + loads.append(model) + return real_select_tokenizer(model="gpt-5.4-mini") if len(loads) == 1 else hub_tokenizer + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hub_up_on_the_second_load) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", reuse_seconds) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 1) + request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]) + + first: Final = await token_counter(request=request, call_endpoint=True) + second: Final = await token_counter(request=request, call_endpoint=True) + + assert (first.tokenizer_type, second.tokenizer_type) == tokenizer_types + + +async def _fork_guard_traceback_frames() -> int: + with pytest.raises(ForkedAfterNativeRuntimeStarted) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + return len(traceback.extract_tb(raised.value.__traceback__)) + + +@pytest.mark.asyncio +async def test_fork_guard_error_while_loading_the_tokenizer_is_raised_to_every_count_sharing_the_load( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi(monkeypatch) + + def forked_load(model: str, custom_tokenizer: object = None) -> NoReturn: + raise ForkedAfterNativeRuntimeStarted("native runtime started before fork") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", forked_load) + + frames: Final = [await _fork_guard_traceback_frames() for _ in range(3)] + + assert len(set(frames)) == 1 + + +@pytest.mark.asyncio +async def test_swallowed_estimate_error_is_logged_by_type_without_the_request_content( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + verbose_proxy_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name), pytest.raises(ProxyException): + await token_counter( + request=TokenCountRequest( + model="kimi", + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok", "tool_calls": "private-marker-in-the-prompt"}, + ], + ), + call_endpoint=True, + ) + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + warnings: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert any("count error 400" in message and "ValueError" in message for message in warnings) + assert not any("private-marker-in-the-prompt" in message for message in warnings) + + +@pytest.mark.asyncio +async def test_fork_guard_error_in_the_estimate_is_never_answered_as_the_provider_error(): + async def forked_estimate() -> NoReturn: + raise ForkedAfterNativeRuntimeStarted("native runtime started before fork") + + with pytest.raises(ForkedAfterNativeRuntimeStarted): + await _estimate_or_provider_error( + forked_estimate(), provider_failure=ProviderCountFailure(status_code=400, message="scripted 400") + ) + + +@pytest.mark.asyncio +async def test_local_estimate_adds_the_configured_framing_once_per_request(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": f"sk-unused-{name}"}, + "model_info": model_info, + } + + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[deployment("framed", {"token_count_framing": 100}), deployment("plain", {})]), + ) + messages: Final = [{"role": "system", "content": "be brief"}, {"role": "user", "content": "hello there"}] + + plain: Final = await token_counter(request=TokenCountRequest(model="plain", messages=messages)) + framed: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + framed_prompt: Final = await token_counter(request=TokenCountRequest(model="framed", prompt="hello there")) + + assert plain.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + assert framed.total_tokens == plain.total_tokens + 100 + assert (plain.count_source, framed.count_source) == ("local_estimate", "local_estimate") + assert framed_prompt.total_tokens == litellm.token_counter(model="gpt-5.4-mini", text="hello there") + + +@pytest.mark.asyncio +async def test_framing_comes_from_the_row_the_deployment_provider_or_base_model_names(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, litellm_params: dict[str, object], model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {**litellm_params, "api_key": "unused", "api_base": "https://unused.invalid"}, + "model_info": model_info, + } + + kimi: Final = "databricks/system.ai.kimi-k3" + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + deployment("prefixed", {"model": kimi}, {}), + deployment("split", {"model": "system.ai.kimi-k3", "custom_llm_provider": "databricks"}, {}), + deployment("aliased", {"model": "gpt-5.4-mini", "custom_llm_provider": "openai"}, {"base_model": kimi}), + ] + ), + ) + messages: Final = [{"role": "user", "content": "hello there"}] + framing: Final = cost_map_entry(kimi)["token_count_framing"] + + prefixed: Final = await token_counter(request=TokenCountRequest(model="prefixed", messages=messages)) + split: Final = await token_counter(request=TokenCountRequest(model="split", messages=messages)) + aliased: Final = await token_counter(request=TokenCountRequest(model="aliased", messages=messages)) + + assert prefixed.total_tokens == litellm.token_counter(model="system.ai.kimi-k3", messages=messages) + framing + assert split.total_tokens == prefixed.total_tokens + assert aliased.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + framing + + +_KIMI_WEATHER: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +def _byte_level_kimi_k3(path: Path) -> KimiK3TokenizerSource: + path.write_bytes(b"\n".join(base64.b64encode(bytes((value,))) + b" %d" % value for value in range(256))) + return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + + +def _available_kimi_k3(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> tiktoken.Encoding: + source: Final = _byte_level_kimi_k3(tmp_path / "tiktoken.model") + monkeypatch.setattr( + litellm.proxy.proxy_server, "_KIMI_K3_ENCODINGS", KimiK3Encodings(source=source, retry_after_seconds=60) + ) + encoding: Final = load_kimi_k3_encoding(source) + assert encoding is not None + return encoding + + +def _shipped_kimi_k3_framing() -> int: + framing: Final = cost_map_entry("bedrock/us.moonshotai.kimi-k3")["token_count_framing"] + assert isinstance(framing, int) + return framing + + +@pytest.mark.asyncio +async def test_bedrock_kimi_k3_is_counted_with_its_tokenizer_plus_the_shipped_framing_in_both_request_formats( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + _framed_kimi(monkeypatch) + encoding: Final = _available_kimi_k3(monkeypatch, tmp_path) + question: Final = {"role": "user", "content": "Weather in Paris?"} + openai_messages: Final = [ + {"role": "system", "content": "Be terse."}, + question, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}} + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ] + + openai: Final = await count_request_tokens( + request=TokenCountRequest(model="kimi", messages=openai_messages, tools=[_KIMI_WEATHER]), call_endpoint=True + ) + anthropic: Final = await count_request_tokens( + request=TokenCountRequest( + model="kimi", + system="Be terse.", + messages=[ + question, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "c1", "name": "get_weather", "input": {"city": "Paris"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "Sun"}]}, + ], + tools=[ + { + "name": "get_weather", + "description": "Get the weather", + "input_schema": _KIMI_WEATHER["function"]["parameters"], + } + ], + ), + call_endpoint=True, + messages_format="anthropic", + ) + exact: Final = kimi_k3_tokens(encoding, "us.moonshotai.kimi-k3", openai_messages, [_KIMI_WEATHER]) + + assert exact is not None + assert (openai.total_tokens, openai.tokenizer_type, openai.count_source) == ( + exact + _shipped_kimi_k3_framing(), + "huggingface_tokenizer", + "local_estimate", + ) + assert (anthropic.total_tokens, anthropic.tokenizer_type) == (openai.total_tokens, "huggingface_tokenizer") + + +@pytest.mark.asyncio +async def test_bedrock_kimi_k3_deployment_with_a_custom_tokenizer_keeps_counting_with_it( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/kimi-custom-tokenizer") + _available_kimi_k3(monkeypatch, tmp_path) + custom: Final = litellm.utils._select_tokenizer(model="gpt-5.4-mini") + monkeypatch.setattr(litellm.utils, "_select_tokenizer", lambda model, custom_tokenizer=None: custom) + messages: Final = [{"role": "user", "content": "What is new?"}] + + counted: Final = await token_counter(request=TokenCountRequest(model="kimi", messages=messages), call_endpoint=True) + + assert counted.total_tokens == ( + litellm.token_counter(model=counted.model_used, messages=messages, custom_tokenizer=custom) + + kimi_k3_template_tokens(len(messages), False) + + _shipped_kimi_k3_framing() + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("messages", "tools", "template_tokens"), + [ + pytest.param([{"role": "user", "content": "What is new?"}], None, 14, id="one_message"), + pytest.param( + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": "Sunny."}, + {"role": "user", "content": "And Oslo?"}, + ], + [_KIMI_WEATHER], + 14 + 3 * 13 + 40, + id="turns_and_tools", + ), + ], +) +async def test_bedrock_kimi_k3_without_its_tokenizer_adds_the_template_offsets_to_the_local_estimate( + monkeypatch: pytest.MonkeyPatch, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + template_tokens: int, +): + _framed_kimi(monkeypatch) + + counted: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=messages, tools=tools), call_endpoint=True + ) + estimate: Final = litellm.token_counter( + model=counted.model_used, + messages=messages, + tools=tools, + custom_tokenizer=litellm.utils._select_tokenizer(model=counted.model_used), + ) + + assert counted.total_tokens == estimate + template_tokens + _shipped_kimi_k3_framing() + + +@pytest.mark.asyncio +async def test_slow_cost_map_lookup_does_not_stall_other_requests_on_the_worker(monkeypatch: pytest.MonkeyPatch): + real_get_model_info: Final = litellm.get_model_info + lookup_started: Final = threading.Event() + other_request_ran: Final = threading.Event() + released_by_the_other_request: Final[list[bool]] = [] + + def slow_get_model_info(model: str, *args: object, **kwargs: object): + if model == "slow-lookup-model": + lookup_started.set() + released_by_the_other_request.append(other_request_ran.wait(5)) + return real_get_model_info(model, *args, **kwargs) + + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "get_model_info", slow_get_model_info) + counting: Final = asyncio.create_task( + token_counter( + request=TokenCountRequest(model="slow-lookup-model", messages=[{"role": "user", "content": "hi"}]) + ) + ) + await asyncio.to_thread(lookup_started.wait, 5) + other_request_ran.set() + await counting + + assert set(released_by_the_other_request) == {True} + + +@pytest.mark.asyncio +async def test_uncountable_model_keeps_the_provider_error_when_local_counting_is_disabled( + monkeypatch: pytest.MonkeyPatch, +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "kimi", + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": {"token_count_framing": 78}, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + monkeypatch.setattr(litellm, "disable_token_counter", True) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert raised.value.code == "400" + assert "scripted 400 for kimi" in raised.value.message + + +@pytest.mark.parametrize( + ("model_info", "counted_with"), + [ + ({}, ("us.anthropic.claude-opus-5-5", "us-east-1", 447)), + ( + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}}, + ("us.anthropic.claude-opus-5-5", "ap-southeast-4", 447), + ), + ({"token_count_model": {}}, ("us.anthropic.claude-sonnet-5-5", "us-west-2", 300)), + ], +) +@pytest.mark.usefixtures("bedrock_sonnet_5_5_mapped_as_on_main") +@pytest.mark.asyncio +async def test_configured_token_count_model_overrides_the_shipped_one_and_an_empty_one_counts_the_model_itself( + monkeypatch: pytest.MonkeyPatch, model_info: dict[str, object], counted_with: tuple[str, str, int] +): + counter: Final = _ScriptedCounter( + { + "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=300), + "us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447), + } + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-sonnet-5-5", + "aws_region_name": "us-west-2", + }, + "model_info": model_info, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + counted: Final = await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert (counter.calls[0][0], counter.regions[0], counted.total_tokens) == counted_with + assert len(counter.calls) == 1 + + +@pytest.mark.parametrize( + ("aws_region_name", "counted_in"), + [("eu-central-1", "eu-central-1"), ("us-west-2", "us-east-1")], +) +@pytest.mark.usefixtures("bedrock_sonnet_5_5_mapped_as_on_main") +@pytest.mark.asyncio +async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geography( + monkeypatch: pytest.MonkeyPatch, aws_region_name: str, counted_in: str +): + counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": { + "model": "bedrock/anthropic.claude-sonnet-5-5", + "aws_region_name": aws_region_name, + }, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert (counter.calls[0][0], counter.regions[0]) == ("anthropic.claude-opus-5-5", counted_in) + + +class _SlowVault: + def sync_read_secret(self, secret_name: str) -> None: + time.sleep(0.5) + + +async def _longest_stall(done: asyncio.Event) -> float: + if done.is_set(): + return 0.0 + started: Final = time.monotonic() + await asyncio.sleep(0.01) + return max(time.monotonic() - started, await _longest_stall(done)) + + +@pytest.mark.usefixtures("bedrock_sonnet_5_5_mapped_as_on_main") +@pytest.mark.asyncio +async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_free(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + + counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + + def bedrock_counter(deployment: object, model: str) -> tuple[_ScriptedCounter, str, str]: + return counter, model, "bedrock" + + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[{"model_name": "sonnet", "litellm_params": {"model": "bedrock/anthropic.claude-sonnet-5-5"}}] + ), + ) + monkeypatch.setattr(proxy_server, "_get_provider_token_counter", bedrock_counter) + monkeypatch.setattr(litellm, "secret_manager_client", _SlowVault()) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.HASHICORP_VAULT) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + done: Final = asyncio.Event() + + async def count() -> TokenCountResponse: + try: + return await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), + call_endpoint=True, + ) + finally: + done.set() + + started: Final = time.monotonic() + stall, response = await asyncio.gather(_longest_stall(done), count()) + + assert time.monotonic() - started >= 1.0 + assert stall < 0.4 + assert response.total_tokens == 447 + + +@pytest.mark.asyncio +async def test_framing_still_applies_while_the_deployment_is_cooling_down(monkeypatch: pytest.MonkeyPatch): + router: Final = Router( + model_list=[ + { + "model_name": "framed", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-unused-framed"}, + "model_info": {"id": "framed-deployment", "token_count_framing": 100}, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + messages: Final = [{"role": "user", "content": "hello there"}] + warm: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + + router.cooldown_cache.add_deployment_to_cooldown( + model_id="framed-deployment", + original_exception=Exception("rate limited"), + exception_status=429, + cooldown_time=60, + ) + cooling: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + + assert warm.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + 100 + assert cooling.total_tokens - 100 == litellm.token_counter(model="framed", messages=messages) + assert cooling.count_source == "local_estimate" + + class _GeminiPart(BaseModel): model_config = ConfigDict(frozen=True) text: str = "" diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index d73f5efa96b..6a6e4c5b4d4 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -513,6 +513,40 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing(): ) +def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None: + """A deployment's ``token_count_model`` and ``token_count_framing`` must not land on + the shared ``{provider}/{model}`` key, or a sibling deployment of that model would count with them.""" + from litellm.proxy.common_utils.token_count_settings import cost_map_entry + + backend_model: Final = "bedrock/us.moonshotai.kimi-k3" + deploy_id: Final = "kimi-deploy-keep-token-count" + configured: Final = { + "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-west-2"}, + "token_count_framing": 7, + } + builtin: Final = {field: cost_map_entry(backend_model).get(field) for field in configured} + assert all(builtin[field] != value for field, value in configured.items()), builtin + model_keys: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in (deploy_id, backend_model, backend_model.removeprefix("bedrock/")) + } + try: + Router( + model_list=[ + { + "model_name": "kimi-keep-token-count", + "litellm_params": {"model": backend_model}, + "model_info": {"id": deploy_id, **configured}, + } + ], + ) + + assert {field: litellm.model_cost[deploy_id][field] for field in configured} == configured + assert {field: cost_map_entry(backend_model).get(field) for field in configured} == builtin + finally: + _restore_model_cost_entries(model_keys) + + def test_should_store_full_pricing_under_deployment_model_id(): """ Per-deployment pricing (including zero) should be stored and diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a934f57f3d3..e80ca2a5820 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -936,6 +936,13 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_video_input": {"type": "boolean"}, "supports_vision": {"type": "boolean"}, "supports_web_search": {"type": "boolean"}, + "token_count_framing": {"type": "integer"}, + "token_count_model": { + "type": "object", + "properties": {"model": {"type": "string"}, "region": {"type": "string"}}, + "required": ["model"], + }, + "token_count_tokenizer": {"type": "string", "enum": ["kimi_k3"]}, "supports_url_context": {"type": "boolean"}, "supports_multimodal": {"type": "boolean"}, "uses_embed_content": {"type": "boolean"}, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 16eaa1e3630..42c306cc3b3 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -43384,9 +43384,17 @@ export interface components { * @description Response structure for token count details with modality breakdown. * * Example: - * {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]} + * {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}], + * 'litellm_count_source': 'provider'} */ TokenCountDetailsResponse: { + /** Litellm Count Model */ + litellm_count_model?: string; + /** + * Litellm Count Source + * @enum {string} + */ + litellm_count_source?: "provider" | "token_count_model" | "local_estimate"; /** Prompttokensdetails */ promptTokensDetails: components["schemas"]["PromptTokensDetails"][]; /** Totaltokens */ @@ -43415,6 +43423,8 @@ export interface components { }; /** TokenCountResponse */ TokenCountResponse: { + /** Count Source */ + count_source?: ("provider" | "token_count_model" | "local_estimate") | null; /** * Error * @default false From c3d07e70d6e60dc7997e1777fcb35a35444be96c Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:17:01 -0700 Subject: [PATCH 24/36] fix(count_tokens)!: count tools the way DeepSeek V4 and Qwen3.5 chat templates render them (#47173) Backport of #47173 to stable/1.104.x. Cherry-picked from c82e8244a57862feecfb48218bdd20c72c2afe6d (main). --- .../prompt_templates/chat_template_tools.py | 37 ++ litellm/litellm_core_utils/token_counter.py | 114 ++++-- .../test_chat_template_tool_count_wire.py | 383 ++++++++++++++++++ ...ls_prompt_templates_chat_template_tools.py | 31 ++ .../litellm_core_utils/test_token_counter.py | 181 ++++++++- 5 files changed, 721 insertions(+), 25 deletions(-) create mode 100644 litellm/litellm_core_utils/prompt_templates/chat_template_tools.py create mode 100644 tests/integration/providers/test_chat_template_tool_count_wire.py create mode 100644 tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py diff --git a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py new file mode 100644 index 00000000000..3c437ca7fbb --- /dev/null +++ b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py @@ -0,0 +1,37 @@ +import json +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Final + + +@dataclass(frozen=True, slots=True) +class ChatTemplateToolSection: + preamble_tokens: int + implicit_system_turn_tokens: int + render_tool_schemas: Callable[[Sequence[Mapping[str, object]]], str] + + +def _deepseek_v4_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: + return "\n".join(json.dumps(function, ensure_ascii=False) for function in functions) + + +def _qwen3_5_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: + return "".join( + "\n" + json.dumps({"type": "function", "function": function}, ensure_ascii=False) for function in functions + ) + + +_TOOL_SECTIONS: Final = ( + ( + re.compile(r"deepseek-v4(?![.\d]|p\d)", re.IGNORECASE), + ChatTemplateToolSection(222, 0, _deepseek_v4_tool_schemas), + ), + (re.compile(r"qwen3[._]?5(?!\d)", re.IGNORECASE), ChatTemplateToolSection(196, 5, _qwen3_5_tool_schemas)), +) + + +def chat_template_tool_section(model: str | None) -> ChatTemplateToolSection | None: + if not model: + return None + return next((section for pattern, section in _TOOL_SECTIONS if pattern.search(model)), None) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index b4a0d7cca21..5a76e85da75 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -34,6 +34,10 @@ from litellm.constants import ( TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.prompt_templates.chat_template_tools import ( + ChatTemplateToolSection, + chat_template_tool_section, +) from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace, HuggingFaceTokenizer, OpenAIEncoding from litellm.litellm_core_utils.url_utils import safe_get from litellm.llms.custom_httpx.http_handler import _get_httpx_client @@ -343,6 +347,9 @@ EXTRAPOLATION_SAMPLES: Final = 16 T_ParamSpec: Final = ParamSpec("T_ParamSpec") T_Retval = TypeVar("T_Retval") _COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter") +_ANTHROPIC_TOOL_KEYS_MAPPED_BY_BRIDGE: Final = frozenset( + ("name", "input_schema", "description", "cache_control", "strict", "type", "eager_input_streaming") +) def _count_offload_limiter_for_this_loop() -> anyio.CapacityLimiter: @@ -400,6 +407,7 @@ class _MessageCountParams: self.tokens_per_message = 3 self.tokens_per_name = 1 self.count_function = _get_count_function(model, custom_tokenizer) + self.tool_section = chat_template_tool_section(model) def token_counter( @@ -463,7 +471,9 @@ def token_counter( num_tokens = _count_messages(params, new_messages, use_default_image_token_count, default_token_count) if count_response_tokens is False: includes_system_message: Final = any(message.get("role", None) == "system" for message in new_messages) - num_tokens += _count_extra(params.count_function, tools, tool_choice, includes_system_message) + num_tokens += _count_extra( + params.count_function, tools, tool_choice, includes_system_message, params.tool_section + ) else: raise ValueError("Either text or messages must be provided") @@ -557,6 +567,7 @@ def _count_extra( tools: list[ChatCompletionToolParam] | None, tool_choice: ChatCompletionNamedToolChoiceParam | None, includes_system_message: bool, + tool_section: ChatTemplateToolSection | None, ) -> int: """Count extra tokens for function definitions and tool choices. Args: @@ -564,26 +575,56 @@ def _count_extra( tools (Optional[List[ChatCompletionToolParam]]): The available tools. tool_choice (Optional[ChatCompletionNamedToolChoiceParam]): The tool choice. includes_system_message (bool): Whether the messages include a system message. + tool_section (Optional[ChatTemplateToolSection]): How the model's chat template renders tools, if known. """ + reply_primer_tokens: Final = 3 + tool_choice_tokens: Final = 0 if tool_section is not None else _count_tool_choice(count_function, tool_choice) + return ( + reply_primer_tokens + + _count_tool_definitions(count_function, tools, includes_system_message, tool_section) + + tool_choice_tokens + ) - num_tokens = 3 # every reply is primed with <|start|>assistant<|message|> - if tools: - num_tokens += count_function(_format_function_definitions(tools)) - num_tokens += 9 # Additional tokens for function definition of tools - # If there's a system message and tools are present, subtract four tokens - if tools and includes_system_message: - num_tokens -= 4 - # If tool_choice is 'none', add one token. - # If it's an object, add 4 + the number of tokens in the function name. - # If it's undefined or 'auto', don't add anything. +def _count_tool_definitions( + count_function: TokenCounterFunction, + tools: Sequence[object] | None, + includes_system_message: bool, + tool_section: ChatTemplateToolSection | None, +) -> int: + if not tools: + return 0 + if tool_section is not None: + return _count_chat_template_tools(count_function, tools, includes_system_message, tool_section) + system_message_discount: Final = 4 if includes_system_message else 0 + return count_function(_format_function_definitions(tools)) + 9 - system_message_discount + + +def _count_tool_choice( + count_function: TokenCounterFunction, + tool_choice: ChatCompletionNamedToolChoiceParam | None, +) -> int: if tool_choice == "none": - num_tokens += 1 - elif isinstance(tool_choice, dict): - num_tokens += 7 - num_tokens += count_function(str(tool_choice["function"]["name"])) + return 1 + if isinstance(tool_choice, dict): + return 7 + count_function(str(tool_choice["function"]["name"])) + return 0 - return num_tokens + +def _count_chat_template_tools( + count_function: TokenCounterFunction, + tools: Sequence[object], + includes_system_message: bool, + tool_section: ChatTemplateToolSection, +) -> int: + functions: Final = tuple( + function for function in _function_definitions(_tools_sent_as_functions(tools)) if function.get("name") + ) + if not functions: + return 0 + system_turn_tokens: Final = 0 if includes_system_message else tool_section.implicit_system_turn_tokens + schemas: Final = tool_section.render_tool_schemas(functions) + return tool_section.preamble_tokens + system_turn_tokens + count_function(schemas) def _get_extrapolating_count_function( @@ -1000,15 +1041,31 @@ def _format_function_definitions(tools: Sequence[object]) -> str: lines: Final = [] lines.append("namespace functions {") lines.append("") - for tool in tools: - if not isinstance(tool, Mapping): - continue - for function in _function_definitions_for_tool(cast(Mapping[str, object], tool)): - lines.extend(_format_single_function_definition(function)) + for function in _function_definitions(tools): + lines.extend(_format_single_function_definition(function)) lines.append("} // namespace functions") return "\n".join(lines) +def _tools_sent_as_functions(tools: Sequence[object]) -> tuple[Mapping[str, object], ...]: + return tuple(tool for tool in _tool_mappings(tools) if not _is_anthropic_web_search_tool(tool)) + + +def _is_anthropic_web_search_tool(tool: Mapping[str, object]) -> bool: + return tool.get("name") == "web_search" and "input_schema" not in tool and tool.get("type") != "function" + + +def _tool_mappings(tools: Sequence[object]) -> Iterable[Mapping[str, object]]: + for tool in tools: + if isinstance(tool, Mapping): + yield cast(Mapping[str, object], tool) + + +def _function_definitions(tools: Sequence[object]) -> Iterable[Mapping[str, object]]: + for tool in _tool_mappings(tools): + yield from _function_definitions_for_tool(tool) + + def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mapping[str, object]]: function: Final = tool.get("function") if isinstance(function, Mapping): @@ -1020,12 +1077,21 @@ def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mappi if isinstance(declaration, Mapping): yield declaration return - parameters: Final = tool.get("input_schema") or tool.get("parameters") or {} + input_schema: Final = tool.get("input_schema") + parameters: Final = input_schema or tool.get("parameters") or {} normalized_parameters: Final = parameters if isinstance(parameters, Mapping) else {} + unmapped_anthropic_keys: Final = ( + {key: value for key, value in tool.items() if key not in _ANTHROPIC_TOOL_KEYS_MAPPED_BY_BRIDGE} + if input_schema is not None + else {} + ) + description: Final = tool.get("description") + strict: Final = tool.get("strict") yield { "name": tool.get("name"), - "description": tool.get("description"), - "parameters": normalized_parameters, + "parameters": {**normalized_parameters, **unmapped_anthropic_keys}, + **({"description": description} if description is not None else {}), + **({"strict": bool(strict)} if strict is not None else {}), } diff --git a/tests/integration/providers/test_chat_template_tool_count_wire.py b/tests/integration/providers/test_chat_template_tool_count_wire.py new file mode 100644 index 00000000000..d2b85e7168f --- /dev/null +++ b/tests/integration/providers/test_chat_template_tool_count_wire.py @@ -0,0 +1,383 @@ +import json +import re +import uuid +from collections.abc import Iterator, Mapping, Sequence +from contextlib import ExitStack +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.openai_wire import chat_reply +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from litellm.litellm_core_utils.default_encoding import encoding as cl100k_base +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_API_KEY: Final = "scripted-provider-key" +_DEEPSEEK_V4: Final = "azure_ai/DeepSeek-V4-Flash-0731" +_QWEN3_5: Final = "databricks/databricks-qwen35-122b-a10b" +_GROK: Final = "azure_ai/grok-4-1" +_DEEPSEEK_V4_TOOL_PREAMBLE: Final = 222 +_QWEN3_5_TOOL_PREAMBLE: Final = 196 +_QWEN3_5_OPENED_SYSTEM_TURN: Final = 5 +_OPENAI_TOOL_OVERHEAD: Final = 9 +_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], +} +_FUNCTION: Final[dict[str, JsonValue]] = {"name": "lookup", "parameters": _SCHEMA} +_OPENAI_NAMESPACE: Final = "namespace functions {\n\ntype lookup = (_: {\nquery: string,\n}) => any;\n\n} // namespace functions" +_ANTHROPIC_TOOL: Final[dict[str, JsonValue]] = {"name": "lookup", "input_schema": _SCHEMA} +_OPENAI_TOOL: Final[dict[str, JsonValue]] = {"type": "function", "function": _FUNCTION} +_RESPONSES_TOOL: Final[dict[str, JsonValue]] = {"type": "function", **_FUNCTION} +_NAMED_TOOL_CHOICE: Final[dict[str, JsonValue]] = {"type": "function", "function": {"name": "lookup"}} +_WEB_SEARCH_TOOL: Final[dict[str, JsonValue]] = {"type": "web_search_20250305", "name": "web_search", "max_uses": 3} +_PROMPT: Final = "Find record 14." +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": _PROMPT}] +_INPUT_PRICE: Final = 0.001 +_OUTPUT_PRICE: Final = 0.002 +_TINY_BUDGET: Final = 0.000000001 +_ESTIMATED_COST: Final = re.compile(r"Estimated request cost: ([^\s,]+)") + + +def _tokens(text: str) -> int: + return len(cl100k_base.encode(text)) + + +def _deepseek_v4_tool_tokens(function: Mapping[str, JsonValue]) -> int: + return _DEEPSEEK_V4_TOOL_PREAMBLE + _tokens(json.dumps(function, ensure_ascii=False)) + + +def _qwen3_5_tool_tokens(function: Mapping[str, JsonValue], *, with_system: bool) -> int: + opened_system_turn: Final = 0 if with_system else _QWEN3_5_OPENED_SYSTEM_TURN + schema: Final = json.dumps({"type": "function", "function": function}, ensure_ascii=False) + return _QWEN3_5_TOOL_PREAMBLE + opened_system_turn + _tokens("\n" + schema) + + +def _completion(_request: Request) -> Reply: + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", "scripted-model", "OK", stream=False) + + +def _deployment(scenario: Scenario, group: str, model: str, api_base: str) -> str: + name: Final = f"{group}-{uuid.uuid4().hex}" + created: Final = scenario.gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": {"model": model, "api_base": api_base, "api_key": _API_KEY}, + "model_info": {}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return name + + +@dataclass(frozen=True, slots=True) +class _CountRoute: + path: str + field: str + tool: dict[str, JsonValue] + + def body(self, model: str, tools: Sequence[JsonValue]) -> dict[str, JsonValue]: + prompt: Final[dict[str, JsonValue]] = ( + {"input": _PROMPT} if self.path == "/v1/responses/input_tokens" else {"messages": _MESSAGES} + ) + return {"model": model, **prompt, **({"tools": list(tools)} if tools else {})} + + +_COUNT_ROUTES: Final[Mapping[str, _CountRoute]] = { + "messages-count-tokens": _CountRoute("/v1/messages/count_tokens", "input_tokens", _ANTHROPIC_TOOL), + "utils-token-counter": _CountRoute("/utils/token_counter", "total_tokens", _OPENAI_TOOL), + "responses-input-tokens": _CountRoute("/v1/responses/input_tokens", "input_tokens", _RESPONSES_TOOL), +} +_MESSAGES_COUNT: Final = _COUNT_ROUTES["messages-count-tokens"] + + +def _count(gateway: Gateway, route: _CountRoute, body: Mapping[str, JsonValue]) -> int: + response: Final = gateway.request("POST", route.path, body) + assert response.status_code == 200, response.text + counted: Final = JSON_OBJECT.validate_json(response.content)[route.field] + assert isinstance(counted, int), response.text + return counted + + +@pytest.mark.parametrize("route", sorted(_COUNT_ROUTES)) +def test_count_routes_count_deepseek_v4_tools_the_way_its_chat_template_renders_them( + gateway: Gateway, route: str +) -> None: + selected: Final = _COUNT_ROUTES[route] + with wire_server(_completion) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "count", _DEEPSEEK_V4, wire.url) + without_tools: Final = _count(gateway, selected, selected.body(model, ())) + with_tool: Final = _count(gateway, selected, selected.body(model, (selected.tool,))) + upstream: Final = wire.drain() + assert upstream == () + assert with_tool - without_tools == _deepseek_v4_tool_tokens(_FUNCTION) + + +@pytest.mark.parametrize("with_system", (False, True), ids=("no-system", "system")) +def test_messages_count_tokens_counts_qwen3_5_tools_the_way_its_chat_template_renders_them( + gateway: Gateway, with_system: bool +) -> None: + system: Final[dict[str, JsonValue]] = {"system": "Be brief."} if with_system else {} + with wire_server(_completion) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "count", _QWEN3_5, wire.url) + without_tools: Final = _count(gateway, _MESSAGES_COUNT, {**_MESSAGES_COUNT.body(model, ()), **system}) + with_tool: Final = _count( + gateway, _MESSAGES_COUNT, {**_MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,)), **system} + ) + upstream: Final = wire.drain() + assert upstream == () + assert with_tool - without_tools == _qwen3_5_tool_tokens(_FUNCTION, with_system=with_system) + + +def test_messages_count_tokens_skips_web_search_and_counts_unmapped_tool_keys_for_deepseek_v4( + gateway: Gateway, +) -> None: + deferred_function: Final[dict[str, JsonValue]] = {"name": "lookup", "parameters": {**_SCHEMA, "defer_loading": True}} + with wire_server(_completion) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "count", _DEEPSEEK_V4, wire.url) + plain: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,))) + with_web_search: Final = _count( + gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL, _WEB_SEARCH_TOOL)) + ) + deferred: Final = _count( + gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ({**_ANTHROPIC_TOOL, "defer_loading": True},)) + ) + upstream: Final = wire.drain() + assert upstream == () + assert with_web_search == plain + assert deferred - plain == _deepseek_v4_tool_tokens(deferred_function) - _deepseek_v4_tool_tokens(_FUNCTION) + + +def test_messages_count_tokens_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None: + with wire_server(_completion) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, "count", _GROK, wire.url) + without_tools: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ())) + plain: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,))) + deferred: Final = _count( + gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ({**_ANTHROPIC_TOOL, "defer_loading": True},)) + ) + upstream: Final = wire.drain() + assert upstream == () + assert plain - without_tools == _OPENAI_TOOL_OVERHEAD + _tokens(_OPENAI_NAMESPACE) + assert deferred == plain + + +@dataclass(frozen=True, slots=True) +class _Endpoint: + path: str + body: dict[str, JsonValue] + tools: dict[str, JsonValue] + + +_ENDPOINTS: Final[Mapping[str, _Endpoint]] = { + "chat-completions": _Endpoint( + "/v1/chat/completions", + {"messages": _MESSAGES, "max_tokens": 1}, + {"tools": [_OPENAI_TOOL], "tool_choice": _NAMED_TOOL_CHOICE}, + ), + "messages": _Endpoint( + "/v1/messages", + {"messages": _MESSAGES, "max_tokens": 1}, + {"tools": [_ANTHROPIC_TOOL], "tool_choice": {"type": "auto"}}, + ), + "responses": _Endpoint( + "/v1/responses", + {"input": _PROMPT, "max_output_tokens": 16}, + {"tools": [_RESPONSES_TOOL]}, + ), +} + + +def _send(gateway: Gateway, endpoint: _Endpoint, group: str, key: str, *, tools: bool) -> httpx.Response: + return gateway.request( + "POST", + endpoint.path, + {"model": group, **endpoint.body, **(endpoint.tools if tools else {})}, + key=key, + ) + + +@dataclass(frozen=True, slots=True) +class _BudgetRig: + gateway: Gateway + wire: Wire + group: str + + +def _strict_budget_config(directory: Path, group: str, api_base: str) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "model_list": [ + { + "model_name": group, + "litellm_params": { + "model": _DEEPSEEK_V4, + "api_base": api_base, + "api_key": _API_KEY, + "input_cost_per_token": _INPUT_PRICE, + "output_cost_per_token": _OUTPUT_PRICE, + }, + } + ], + "general_settings": {**object_value(base["general_settings"]), "fail_closed_budget_enforcement": True}, + } + path: Final = directory / f"chat-template-tool-budget-{uuid.uuid4().hex[:8]}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def strict_budget(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_BudgetRig]: + directory: Final = tmp_path_factory.mktemp("chat-template-tool-budget") + group: Final = f"deepseek-v4-flash-{uuid.uuid4().hex}" + with ExitStack() as stack: + rig_gateway: Final = stack.enter_context(gateway_from_environment()) + wire: Final = stack.enter_context(wire_server(_completion)) + owned: Final = stack.enter_context( + owned_proxy_process(rig_gateway, directory, {}, config=_strict_budget_config(directory, group, wire.url)) + ) + yield _BudgetRig(owned.gateway, wire, group) + + +def _estimated_cost(response: httpx.Response) -> float: + assert response.status_code in (400, 422, 429), f"{response.status_code} {response.text}" + error: Final = object_value(JSON_OBJECT.validate_json(response.content)["error"]) + estimate: Final = _ESTIMATED_COST.search(string_value(error["message"])) + assert estimate is not None, response.text + return float(estimate[1]) + + +@pytest.mark.parametrize("endpoint", ("chat-completions", "messages")) +def test_budget_reservation_estimates_deepseek_v4_tools_the_way_its_chat_template_renders_them( + strict_budget: _BudgetRig, endpoint: str +) -> None: + selected: Final = _ENDPOINTS[endpoint] + with strict_budget.gateway.scenario() as scenario: + key: Final = scenario.key(models=[strict_budget.group], max_budget=_TINY_BUDGET) + without_tools: Final = _estimated_cost(_send(strict_budget.gateway, selected, strict_budget.group, key, tools=False)) + with_tools: Final = _estimated_cost(_send(strict_budget.gateway, selected, strict_budget.group, key, tools=True)) + assert strict_budget.wire.drain() == () + assert with_tools - without_tools == pytest.approx(_deepseek_v4_tool_tokens(_FUNCTION) * _INPUT_PRICE) + + +def _project_key(scenario: Scenario, group: str) -> str: + team: Final = scenario.team() + project: Final = scenario.project(team, model_itpm_limit={group: _deepseek_v4_tool_tokens(_FUNCTION)}) + return scenario.key(team_id=team, project_id=project) + + +@pytest.mark.parametrize("endpoint", sorted(_ENDPOINTS)) +def test_project_itpm_reserves_deepseek_v4_tools_the_way_its_chat_template_renders_them( + gateway: Gateway, endpoint: str +) -> None: + with wire_server(_completion) as wire, gateway.scenario() as scenario: + group: Final = _deployment(scenario, "deepseek-v4-flash", _DEEPSEEK_V4, wire.url) + response: Final = _send(gateway, _ENDPOINTS[endpoint], group, _project_key(scenario, group), tools=True) + upstream: Final = wire.drain() + assert response.status_code == 429, response.text + assert "model_per_project_itpm" in response.text, response.text + assert upstream == () + + +def test_project_itpm_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None: + chat: Final = _ENDPOINTS["chat-completions"] + with wire_server(_completion) as wire, gateway.scenario() as scenario: + group: Final = _deployment(scenario, "grok-4-1", _GROK, wire.url) + response: Final = gateway.request( + "POST", + chat.path, + {"model": group, **chat.body, "tools": [_OPENAI_TOOL]}, + key=_project_key(scenario, group), + ) + upstream: Final = wire.drain() + assert response.status_code == 200, response.text + assert len(upstream) == 1, upstream + + +_FILES: Final = 3 +_READ_FILE_FUNCTION: Final[dict[str, JsonValue]] = { + "name": "read_file", + "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}, +} +_READ_FILE_TOOL: Final[dict[str, JsonValue]] = { + "name": "read_file", + "input_schema": _READ_FILE_FUNCTION["parameters"], +} + + +def _read_file_round(marker: str, index: int) -> tuple[JsonValue, ...]: + return ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}.txt"}} + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": f"toolu_{index}", + "content": f"short file body.", + } + ], + }, + ) + + +def _agent_body(group: str, marker: str) -> dict[str, JsonValue]: + rounds: Final = tuple(chain.from_iterable(_read_file_round(marker, index) for index in range(_FILES))) + return { + "model": group, + "max_tokens": 16, + "tools": [_READ_FILE_TOOL], + "messages": [*rounds, {"role": "user", "content": "Summarize what you read."}], + "context_management": { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": _deepseek_v4_tool_tokens(_READ_FILE_FUNCTION)}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] + }, + } + + +@pytest.mark.parametrize( + ("model", "files_kept"), + ((_DEEPSEEK_V4, (False, False, True)), (_GROK, (True, True, True))), + ids=("deepseek-v4", "grok-4-1"), +) +def test_messages_context_management_trigger_counts_tools_the_way_the_models_chat_template_renders_them( + gateway: Gateway, model: str, files_kept: tuple[bool, ...] +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_completion) as wire, gateway.scenario() as scenario: + group: Final = _deployment(scenario, "context", model, wire.url) + response: Final = gateway.request("POST", "/v1/messages", _agent_body(group, marker)) + upstream: Final = wire.drain() + assert response.status_code == 200, response.text + assert len(upstream) == 1, upstream + sent: Final = upstream[0].body.decode() + assert tuple(f"" in sent for index in range(_FILES)) == files_kept diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py new file mode 100644 index 00000000000..51249e01d4d --- /dev/null +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py @@ -0,0 +1,31 @@ +from typing import Final + +import pytest + +from litellm.litellm_core_utils.prompt_templates.chat_template_tools import chat_template_tool_section + +_DESCRIBED_FUNCTION: Final = {"name": "f", "description": "é", "parameters": {}} +_BARE_FUNCTION: Final = {"name": "g", "parameters": {}} + + +# Tool schemas as rendered by render_tools in encoding/encoding_dsv4.py (deepseek-ai/DeepSeek-V4-Flash) and by +# chat_template.jinja (Qwen/Qwen3.5-122B-A10B), both read on Hugging Face on 2026-10-17 +@pytest.mark.parametrize( + ("model", "rendered"), + [ + ( + "DeepSeek-V4-Flash-0731", + '{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}', + ), + ( + "databricks-qwen35-122b-a10b", + '\n{"type": "function", "function": {"name": "f", "description": "é", "parameters": {}}}' + '\n{"type": "function", "function": {"name": "g", "parameters": {}}}', + ), + ], +) +def test_chat_template_tool_section_renders_schemas_like_the_models_template(model: str, rendered: str): + section: Final = chat_template_tool_section(model) + + assert section is not None + assert section.render_tool_schemas((_DESCRIBED_FUNCTION, _BARE_FUNCTION)) == rendered diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 42c25a319a2..8d2884756cb 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -9,8 +9,9 @@ import subprocess import sys import threading import time -from collections.abc import Mapping +from collections.abc import Callable, Mapping from concurrent.futures import Future, wait +from copy import deepcopy from itertools import accumulate, chain from pathlib import Path from typing import Final @@ -39,6 +40,7 @@ from litellm.litellm_core_utils.token_counter import ( offload_token_count, ) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new +from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter from tests.large_text import text from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, @@ -515,6 +517,183 @@ def test_token_counter_skips_non_mapping_tools(): assert mixed_tokens == valid_tokens +_WEATHER_FUNCTION: Final = { + "name": "get_weather", + "description": "Get the current weather for a city", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string", "description": "City name"}}, + "required": ["location"], + }, +} +_TIME_FUNCTION: Final = { + "name": "get_time", + "description": "Get the local time for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, +} +_WEB_SEARCH_FUNCTION: Final = { + "name": "web_search", + "description": "Search the web", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}, +} + + +def _openai_tool(function: Mapping[str, object]) -> dict[str, object]: + return {"type": "function", "function": function} + + +def _anthropic_tool(function: Mapping[str, object]) -> dict[str, object]: + return {"name": function["name"], "description": function["description"], "input_schema": function["parameters"]} + + +_USER_MESSAGES: Final = ({"role": "user", "content": "What is the weather in Paris today?"},) +_SYSTEM_AND_USER_MESSAGES: Final = ({"role": "system", "content": "You are a helpful assistant."}, *_USER_MESSAGES) + + +@pytest.mark.parametrize("tool_shape", [_openai_tool, _anthropic_tool]) +@pytest.mark.parametrize( + ("model", "messages", "functions", "billed_prompt_tokens"), + [ + # usage.prompt_tokens billed for this request by Azure AI Foundry DeepSeek-V4-Flash-0731, + # Databricks databricks-qwen35-122b-a10b and DashScope qwen3.5-plus on 2026-10-17 + ("DeepSeek-V4-Flash-0731", _USER_MESSAGES, (_WEATHER_FUNCTION,), 295), + ("DeepSeek-V4-Flash-0731", _USER_MESSAGES, (_WEATHER_FUNCTION, _TIME_FUNCTION), 348), + ("DeepSeek-V4-Flash-0731", _SYSTEM_AND_USER_MESSAGES, (_WEATHER_FUNCTION,), 300), + ("databricks-qwen35-122b-a10b", _USER_MESSAGES, (_WEATHER_FUNCTION,), 289), + ("databricks-qwen35-122b-a10b", _USER_MESSAGES, (_WEATHER_FUNCTION, _TIME_FUNCTION), 352), + ("databricks-qwen35-122b-a10b", _SYSTEM_AND_USER_MESSAGES, (_WEATHER_FUNCTION,), 296), + ("qwen3.5-plus", _USER_MESSAGES, (_WEATHER_FUNCTION,), 284), + ("qwen3.5-plus", _USER_MESSAGES, (_WEATHER_FUNCTION, _TIME_FUNCTION), 342), + ], +) +def test_token_counter_tools_match_billed_prompt_for_json_tool_chat_templates( + tool_shape: Callable[[Mapping[str, object]], dict[str, object]], + model: str, + messages: tuple[dict[str, str], ...], + functions: tuple[Mapping[str, object], ...], + billed_prompt_tokens: int, +): + counted: Final = token_counter_new( + model=model, + messages=list(messages), + tools=[tool_shape(function) for function in functions], + ) + + assert abs(counted - billed_prompt_tokens) <= billed_prompt_tokens * 0.05 + + +@pytest.mark.parametrize( + ("model", "counted_like"), + [ + ("DeepSeek-V4-Pro", "DeepSeek-V4-Flash-0731"), + ("accounts/fireworks/models/deepseek-v4-flash", "DeepSeek-V4-Flash-0731"), + ("DeepSeek-V4.1-Flash", "deepseek-v3.2"), + ("accounts/fireworks/models/deepseek-v4p1-flash", "deepseek-v3.2"), + ("Qwen/Qwen3.5-122B-A10B", "databricks-qwen35-122b-a10b"), + ("qwen3-235b-a22b", "deepseek-v3.2"), + ], +) +def test_token_counter_counts_tools_by_the_models_chat_template_family(model: str, counted_like: str): + tools: Final = [_openai_tool(_WEATHER_FUNCTION), _openai_tool(_TIME_FUNCTION)] + + assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) == token_counter_new( + model=counted_like, messages=list(_USER_MESSAGES), tools=tools + ) + + +@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"]) +@pytest.mark.parametrize( + "builtin_tool", + [ + {"type": "web_search_preview"}, + {"type": "web_search_20250305", "name": "web_search", "max_uses": 5}, + {"name": "web_search"}, + ], +) +def test_token_counter_skips_tools_not_sent_as_functions_for_chat_template_families( + model: str, builtin_tool: dict[str, object] +): + weather_tool: Final = _openai_tool(_WEATHER_FUNCTION) + + assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[builtin_tool]) == token_counter_new( + model=model, messages=list(_USER_MESSAGES) + ) + assert token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=[builtin_tool, weather_tool] + ) == token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[weather_tool]) + + +def _responses_tool(function: Mapping[str, object]) -> dict[str, object]: + return {"type": "function", **function} + + +def _gemini_tool(function: Mapping[str, object]) -> dict[str, object]: + return {"function_declarations": [function]} + + +@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"]) +@pytest.mark.parametrize("function", [_WEATHER_FUNCTION, _WEB_SEARCH_FUNCTION]) +@pytest.mark.parametrize("tool_shape", [_responses_tool, _gemini_tool]) +def test_token_counter_counts_function_tools_like_chat_completions_tools_for_chat_template_families( + model: str, function: Mapping[str, object], tool_shape: Callable[[Mapping[str, object]], dict[str, object]] +): + assert token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=[tool_shape(function)] + ) == token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[_openai_tool(function)]) + + +# Azure AI Foundry DeepSeek-V4-Flash-0731 billed 295 prompt tokens for the weather tool with each of +# these tool_choice values on 2026-10-17, the same as without one +@pytest.mark.parametrize("tool_choice", ["none", "required", {"type": "function", "function": {"name": "get_weather"}}]) +def test_token_counter_tool_choice_adds_nothing_for_deepseek_v4(tool_choice: object): + tools: Final = [_openai_tool(_WEATHER_FUNCTION)] + + assert token_counter_new( + model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools, tool_choice=tool_choice + ) == token_counter_new(model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools) + + +@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"]) +@pytest.mark.parametrize( + "anthropic_tool", + [ + {"name": "get_weather", "input_schema": _WEATHER_FUNCTION["parameters"]}, + _anthropic_tool(_WEATHER_FUNCTION), + {**_anthropic_tool(_WEATHER_FUNCTION), "strict": True}, + {**_anthropic_tool(_WEATHER_FUNCTION), "defer_loading": True, "input_examples": [{"location": "Paris"}]}, + _anthropic_tool(_WEB_SEARCH_FUNCTION), + ], +) +def test_token_counter_counts_anthropic_tools_as_sent_to_chat_template_families( + model: str, anthropic_tool: dict[str, object] +): + sent_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai([deepcopy(anthropic_tool)]) # pyright: ignore[reportArgumentType] # raw Anthropic tool dict + + assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[anthropic_tool]) == token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=sent_tools + ) + + +# Qwen3.5's chat_template.jinja opens a system turn for the tool schemas when the request has none, 5 tokens that a +# request with a system message already pays, while DeepSeek V4's encoding_dsv4.py appends them to the first turn +# either way (both read on Hugging Face on 2026-10-17) +@pytest.mark.parametrize( + ("model", "implicit_system_turn_tokens"), [("DeepSeek-V4-Flash-0731", 0), ("databricks-qwen35-122b-a10b", 5)] +) +def test_token_counter_counts_the_tools_system_turn_only_without_a_system_message( + model: str, implicit_system_turn_tokens: int +): + tools: Final = [_openai_tool(_WEATHER_FUNCTION)] + system_cost_with_tools: Final = token_counter_new( + model=model, messages=list(_SYSTEM_AND_USER_MESSAGES), tools=tools + ) - token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) + system_cost_without_tools: Final = token_counter_new( + model=model, messages=list(_SYSTEM_AND_USER_MESSAGES) + ) - token_counter_new(model=model, messages=list(_USER_MESSAGES)) + + assert system_cost_with_tools == system_cost_without_tools - implicit_system_turn_tokens + + class NeedsToleranceUpdateError(Exception): """Custom exception to mark tests that have improved""" From 94828be8b0308f76408594dd98362ca551b43c10 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:17:06 -0700 Subject: [PATCH 25/36] test(integration): add the openai_wire helper the backported chat template wire test imports Hand-ported verbatim to stable/1.104.x from main. --- tests/integration/_support/openai_wire.py | 126 ++++++++++++++++++++++ 1 file changed, 126 insertions(+) create mode 100644 tests/integration/_support/openai_wire.py diff --git a/tests/integration/_support/openai_wire.py b/tests/integration/_support/openai_wire.py new file mode 100644 index 00000000000..bf9d5f194c1 --- /dev/null +++ b/tests/integration/_support/openai_wire.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import json +from collections.abc import Callable +from typing import Final + +from integration._support.wire import Reply, Request, Wire +from pydantic import JsonValue + +from tests.integration._support.provider import MODEL_DISCOVERY + +_USAGE: Final = {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8} + + +def answering_model_discovery(respond: Callable[[Request], Reply]) -> Callable[[Request], Reply]: + def guarded(request: Request) -> Reply: + if (request.method, request.target) == MODEL_DISCOVERY: + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + return respond(request) + + return guarded + + +def posted_targets(wire: Wire) -> tuple[str, ...]: + return tuple(request.target for request in wire.drain() if request.method == "POST") + + +def openai_error(status: int) -> Reply: + return Reply( + status=status, + body=json.dumps({"error": {"message": f"scripted {status}", "type": "server_error", "code": None}}).encode(), + ) + + +def _data_frame(frame: dict[str, JsonValue]) -> bytes: + return b"data: " + json.dumps(frame).encode() + b"\n\n" + + +def _typed_frame(event: dict[str, JsonValue]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def chat_stream(identity: str, model: str, text: str) -> tuple[bytes, bytes, bytes]: + chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": model} + role_only: Final = _data_frame({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}}]}) + content: Final = _data_frame({**chunk, "choices": [{"index": 0, "delta": {"content": text}}]}) + finish: Final = _data_frame( + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": _USAGE} + ) + return (role_only, content, finish + b"data: [DONE]\n\n") + + +def chat_reply(identity: str, model: str, text: str, *, stream: bool) -> Reply: + if stream: + return Reply(content_type="text/event-stream", chunks=chat_stream(identity, model, text)) + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": _USAGE, + } + ).encode() + ) + + +def _response_object(identity: str, model: str, text: str) -> dict[str, JsonValue]: + return { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": model, + "output": [_message_item(identity, text, "completed")], + "usage": {"input_tokens": 5, "output_tokens": 3, "total_tokens": 8}, + } + + +def _message_item(identity: str, text: str, status: str) -> dict[str, JsonValue]: + return { + "id": f"msg_{identity}", + "type": "message", + "role": "assistant", + "status": status, + "content": [{"type": "output_text", "text": text, "annotations": []}] if status == "completed" else [], + } + + +def responses_stream(identity: str, model: str, text: str) -> tuple[bytes, bytes, bytes]: + response: Final = _response_object(identity, model, text) + item: Final = _message_item(identity, text, "in_progress") + opened: Final = _typed_frame( + { + "type": "response.created", + "sequence_number": 0, + "response": {**response, "status": "in_progress", "output": []}, + } + ) + _typed_frame({"type": "response.output_item.added", "sequence_number": 1, "output_index": 0, "item": item}) + delta: Final = _typed_frame( + { + "type": "response.output_text.delta", + "sequence_number": 2, + "item_id": f"msg_{identity}", + "output_index": 0, + "content_index": 0, + "delta": text, + } + ) + closed: Final = _typed_frame( + { + "type": "response.output_item.done", + "sequence_number": 3, + "output_index": 0, + "item": _message_item(identity, text, "completed"), + } + ) + _typed_frame({"type": "response.completed", "sequence_number": 4, "response": response}) + return (opened, delta, closed) + + +def responses_reply(identity: str, model: str, text: str, *, stream: bool) -> Reply: + if stream: + return Reply(content_type="text/event-stream", chunks=responses_stream(identity, model, text)) + return Reply(body=json.dumps(_response_object(identity, model, text)).encode()) From 828ddf0c37710ae3319a124124d5431079225137 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:03:20 -0700 Subject: [PATCH 26/36] fix(count_tokens): narrow stable backport to the October 19 patch Keep #43900, #45301, #46121, #46784, #46759, #46767, #46772, #46815, #47057 and #46777. Port only the Gemini prerequisites from #46430. Remove the deferred Bedrock Mantle behavior and unrelated picks from the final tree. Restore the PDF test provider fixture and deployment, and retain required test imports. --- .../auto_update_price_and_context_window.yml | 5 +- .github/workflows/sync-together-ai-models.yml | 7 +- ci_cd/cost_map_guard.py | 45 +- ci_cd/generate_model_prices_schema.py | 424 +++++ litellm-rust/crates/model-catalog/AGENTS.md | 1 + .../prompt_templates/chat_template_tools.py | 37 - .../mid_conversation_system.py | 16 +- litellm/litellm_core_utils/token_counter.py | 180 +- .../llms/anthropic/count_tokens/handler.py | 3 +- .../anthropic/count_tokens/transformation.py | 70 +- .../anthropic/count_tokens/handler.py | 3 +- .../count_tokens/bedrock_token_counter.py | 181 +- litellm/llms/bedrock/count_tokens/handler.py | 2 +- .../bedrock/count_tokens/mantle_handler.py | 99 -- litellm/llms/bedrock_mantle/common_utils.py | 4 - .../responses/count_tokens/transformation.py | 13 +- litellm/llms/vertex_ai/common_utils.py | 18 - .../count_tokens/handler.py | 36 +- .../common_utils/openai_error_payload.py | 4 - litellm/proxy/proxy_server.py | 4 +- .../proxy/response_api_endpoints/endpoints.py | 97 +- .../_support/bedrock_runtime_peer.py | 278 --- tests/integration/_support/openai_wire.py | 126 -- tests/integration/conftest.py | 30 + .../providers/_count_tokens_system_lift.py | 159 -- ...anthropic_count_tokens_system_lift_wire.py | 581 ------ .../test_bedrock_mantle_count_tokens_wire.py | 1568 ----------------- .../test_chat_template_tool_count_wire.py | 383 ---- .../test_openai_count_tokens_tools_wire.py | 25 - .../test_vertex_gemini_count_tokens_wire.py | 114 +- .../test_vertex_partner_count_tokens_wire.py | 36 - tests/integration/proxy_config.yaml | 7 + .../proxy/google_endpoints/test_endpoints.py | 78 +- .../proxy/proxy_server/test_routes_utils.py | 7 +- .../response_api_endpoints/test_endpoints.py | 214 --- ...ls_prompt_templates_chat_template_tools.py | 31 - ...rompt_templates_mid_conversation_system.py | 21 - .../litellm_core_utils/test_token_counter.py | 253 +-- ...t_anthropic_count_tokens_transformation.py | 291 --- ...est_bedrock_mantle_count_tokens_handler.py | 109 -- .../test_bedrock_token_counter.py | 176 -- ...test_openai_count_tokens_transformation.py | 69 +- .../llms/vertex_ai/count_tokens/__init__.py | 0 .../vertex_ai/test_vertex_ai_common_utils.py | 29 - tests/unit/proxy/test_custom_tokenizer_bug.py | 60 +- tests/unit/proxy/test_proxy_token_counter.py | 90 +- tests/unit/test_cost_map_guard.py | 52 +- ..._window.py => test_model_prices_schema.py} | 198 +++ 48 files changed, 971 insertions(+), 5263 deletions(-) create mode 100644 ci_cd/generate_model_prices_schema.py delete mode 100644 litellm/litellm_core_utils/prompt_templates/chat_template_tools.py delete mode 100644 litellm/llms/bedrock/count_tokens/mantle_handler.py delete mode 100644 tests/integration/_support/bedrock_runtime_peer.py delete mode 100644 tests/integration/_support/openai_wire.py delete mode 100644 tests/integration/providers/_count_tokens_system_lift.py delete mode 100644 tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py delete mode 100644 tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py delete mode 100644 tests/integration/providers/test_chat_template_tool_count_wire.py delete mode 100644 tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py delete mode 100644 tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py delete mode 100644 tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py delete mode 100644 tests/unit/llms/vertex_ai/count_tokens/__init__.py rename tests/unit/{test_model_prices_and_context_window.py => test_model_prices_schema.py} (67%) diff --git a/.github/workflows/auto_update_price_and_context_window.yml b/.github/workflows/auto_update_price_and_context_window.yml index 838ba35f831..7e40a860ee9 100644 --- a/.github/workflows/auto_update_price_and_context_window.yml +++ b/.github/workflows/auto_update_price_and_context_window.yml @@ -24,9 +24,12 @@ jobs: - name: Update JSON Data run: | uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py" + - name: Regenerate JSON Schema + run: | + uv run --frozen python ci_cd/generate_model_prices_schema.py - name: Create Pull Request run: | - git add model_prices_and_context_window.json + git add model_prices_and_context_window.json model_prices_and_context_window.schema.json git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')" gh pr create --title "Update model_prices_and_context_window.json file" \ --body "Automated update for model_prices_and_context_window.json" \ diff --git a/.github/workflows/sync-together-ai-models.yml b/.github/workflows/sync-together-ai-models.yml index 5325510e45e..18d5e3de1eb 100644 --- a/.github/workflows/sync-together-ai-models.yml +++ b/.github/workflows/sync-together-ai-models.yml @@ -41,6 +41,10 @@ jobs: uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md" env: TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }} + - name: Regenerate the JSON schema + if: steps.existing.outputs.open_pr == '' + run: | + uv run --frozen python ci_cd/generate_model_prices_schema.py - name: Create a pull request when the registry changed if: steps.existing.outputs.open_pr == '' run: | @@ -53,7 +57,8 @@ jobs: git config user.email "41898282+github-actions[bot]@users.noreply.github.com" git checkout -b "$branch" git add model_prices_and_context_window.json \ - litellm/model_prices_and_context_window_backup.json + litellm/model_prices_and_context_window_backup.json \ + model_prices_and_context_window.schema.json git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')" gh auth setup-git git push origin "$branch" diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py index 8e34e79a354..14eecf01461 100644 --- a/ci_cd/cost_map_guard.py +++ b/ci_cd/cost_map_guard.py @@ -1,11 +1,11 @@ """Guard the cost map on pull requests. -Every pull request whose diff against its merge base touches one of the two cost map files gets the file -checks: the files parse and the backup copy matches the root file. A pull request that leaves both untouched -skips them, since merging it keeps the base branch's copies and its head tree only carries whatever state the -branch was cut from. Pull requests from the cost map sync bot (branches named litellm_cost_map_sync_*) always -get the file checks and additionally may only touch those two files or the legacy schema -path and may only add or update models. +Every pull request whose diff against its merge base touches one of the three cost map files gets the file +checks: the files parse, the backup copy matches the root file, and the JSON schema is in sync and validates the +map. A pull request that leaves all three untouched skips them, since merging it keeps the base branch's copies +and its head tree only carries whatever state the branch was cut from. Pull requests from the cost map sync bot +(branches named litellm_cost_map_sync_*) always get the file checks and additionally may only touch those three +files and may only add or update models. """ from __future__ import annotations @@ -19,12 +19,13 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final +from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors + COST_MAP_PATH: Final = "model_prices_and_context_window.json" BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json" -LEGACY_SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json" -GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH) +SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json" +GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH) BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_" -SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"}) EVALUATION_OUTPUT_PRICE_SOURCES: Final[Mapping[str, str]] = MappingProxyType({}) CostMap = dict[str, object] @@ -34,6 +35,7 @@ CostMap = dict[str, object] class Snapshot: cost_map: str backup: str + schema: str def _parse_object(text: str, path: str) -> CostMap | str: @@ -44,6 +46,13 @@ def _parse_object(text: str, path: str) -> CostMap | str: return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root" +def _rendered_schema(cost_map: CostMap) -> str: + try: + return render(build_schema(cost_map)) + except SystemExit as error: + return str(error) + + def _evaluation_output_price_failure( key: str, entry: object, evaluation_output_price_sources: Mapping[str, str] ) -> str | None: @@ -69,13 +78,29 @@ def _file_failures( *, evaluation_output_price_sources: Mapping[str, str] = EVALUATION_OUTPUT_PRICE_SOURCES, ) -> tuple[str, ...]: + schema_text: Final = _rendered_schema(head_map) + if not schema_text.startswith("{"): + return (schema_text,) backup_failure: Final = ( () if head.backup == head.cost_map else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",) ) + schema_failure: Final = ( + () + if head.schema == schema_text + else ( + f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; " + "run `python ci_cd/generate_model_prices_schema.py` and commit the result", + ) + ) return ( *backup_failure, + *schema_failure, + *( + f"{COST_MAP_PATH} does not validate against its schema: {error}" + for error in validation_errors(head_map, json.loads(schema_text))[:20] + ), *( failure for key, entry in head_map.items() @@ -110,7 +135,7 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str *( f"bot PRs may only change the cost map files, not {path}" for path in changed_files - if path not in GUARDED_PATHS and path != LEGACY_SCHEMA_PATH + if path not in GUARDED_PATHS ), *(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map), *(f"bot PRs may not remove fields: {ref}" for ref in removed_fields), diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py new file mode 100644 index 00000000000..359df8e6e32 --- /dev/null +++ b/ci_cd/generate_model_prices_schema.py @@ -0,0 +1,424 @@ +from __future__ import annotations + +import json +import sys +from pathlib import Path +from typing import Final, Optional + +import jsonschema + +REPO_ROOT = Path(__file__).parent.parent +PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json" +SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json" + +SPECIAL_ROOT_KEYS = frozenset({"sample_spec", "fallback_generalizations"}) + +JsonSchema = dict + +NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0} +NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0} +BOOLEAN: JsonSchema = {"type": "boolean"} +STRING: JsonSchema = {"type": "string"} +NON_BLANK_STRING: Final[JsonSchema] = {"type": "string", "pattern": r"\S"} +TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"} +WEEKDAY_PATTERN: Final = ( + r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$" +) + +EXTRA_BOOLEAN_KEYS = frozenset( + { + "gemini_native_audio", + "gemini_audio_only_live", + "uses_embed_content", + "use_openai_responses_path", + "bedrock_converse_supports_strict_tools", + "thinking_always_on", + } +) + +HOURS_UTC: Final[JsonSchema] = { + "description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.', + "oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}], +} + +OFF_PEAK_WINDOW: Final[JsonSchema] = { + "type": "object", + "properties": { + "hours_utc": HOURS_UTC, + "weekdays": { + "type": "array", + "description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.", + "items": { + "oneOf": [ + {"type": "integer", "minimum": 1, "maximum": 7}, + {"type": "string", "pattern": WEEKDAY_PATTERN}, + ] + }, + "minItems": 1, + }, + }, + "required": ["hours_utc"], + "additionalProperties": False, +} + +OBJECT_KEYS: dict[str, JsonSchema] = { + "off_peak_pricing": { + "type": "object", + "description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.", + "properties": { + "hours_utc": HOURS_UTC, + "windows": {"type": "array", "items": OFF_PEAK_WINDOW, "minItems": 1}, + "weekday_timezone": { + "type": "string", + "description": "IANA zone the weekdays of each window are read on; defaults to UTC.", + }, + "input_cost_per_token": NONNEG_NUMBER, + "output_cost_per_token": NONNEG_NUMBER, + "output_cost_per_reasoning_token": NONNEG_NUMBER, + "cache_read_input_token_cost": NONNEG_NUMBER, + "cache_creation_input_token_cost": NONNEG_NUMBER, + }, + "anyOf": [{"required": ["hours_utc"]}, {"required": ["windows"]}], + "additionalProperties": False, + }, + "search_context_cost_per_query": { + "type": "object", + "description": "USD cost per web search query, keyed by search context size.", + "properties": { + "search_context_size_low": NONNEG_NUMBER, + "search_context_size_medium": NONNEG_NUMBER, + "search_context_size_high": NONNEG_NUMBER, + }, + "additionalProperties": False, + }, + "guardrail_cost_per_unit": { + "type": "object", + "description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).", + "additionalProperties": NONNEG_NUMBER, + }, + "metadata": { + "type": "object", + "description": "Free-form notes about the entry (e.g. pricing derivation).", + }, + "provider_specific_entry": { + "type": "object", + "description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).", + }, + "token_count_model": { + "type": "object", + "required": ["model"], + "additionalProperties": False, + "properties": { + "model": NON_BLANK_STRING, + "region": NON_BLANK_STRING, + }, + "description": "Model that counts tokens for this one, and the AWS region its count API runs in.", + }, +} + +ARRAY_KEYS: dict[str, JsonSchema] = { + "supported_audio_formats": { + "type": "array", + "description": "Audio container formats the model can return.", + "items": {"type": "string", "enum": ["mp3", "wav"]}, + }, + "supported_endpoints": { + "type": "array", + "description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.", + "items": STRING, + }, + "supported_modalities": { + "type": "array", + "description": "Input modalities the model accepts.", + "items": {"type": "string", "enum": ["text", "image", "audio", "video"]}, + }, + "supported_output_modalities": { + "type": "array", + "description": "Output modalities the model can produce.", + "items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]}, + }, + "reasoning_effort_levels": { + "type": "array", + "description": "Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.", + "items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]}, + }, + "supported_regions": { + "type": "array", + "description": "Cloud regions the model is available in ('global' or region ids).", + "items": STRING, + }, + "tiered_pricing": { + "type": "array", + "description": "Context-length or result-count tiered rates; each tier's costs apply within its range.", + "items": { + "type": "object", + "properties": { + "range": { + "type": "array", + "description": "[min, max] prompt-token span this tier applies to.", + "items": NONNEG_NUMBER, + "minItems": 2, + "maxItems": 2, + }, + "max_results_range": { + "type": "array", + "description": "[min, max] result-count span this tier applies to (search models).", + "items": NONNEG_NUMBER, + "minItems": 2, + "maxItems": 2, + }, + "input_cost_per_token": NONNEG_NUMBER, + "output_cost_per_token": NONNEG_NUMBER, + "output_cost_per_reasoning_token": NONNEG_NUMBER, + "cache_read_input_token_cost": NONNEG_NUMBER, + "cache_creation_input_token_cost": NONNEG_NUMBER, + "input_cost_per_query": NONNEG_NUMBER, + }, + "additionalProperties": False, + }, + }, +} + +INTEGER_KEYS: dict[str, JsonSchema] = { + "max_tokens": { + **NONNEG_INTEGER, + "description": "Legacy field: max output tokens if the provider specifies it, else max input tokens.", + }, + "max_input_tokens": { + **NONNEG_INTEGER, + "description": "Maximum prompt/context tokens the model accepts.", + }, + "max_output_tokens": { + **NONNEG_INTEGER, + "description": "Maximum tokens the model can generate in one response.", + }, + "output_vector_size": { + **NONNEG_INTEGER, + "description": "Embedding dimension for embedding models.", + }, + "prompt_cache_min_tokens": { + **NONNEG_INTEGER, + "description": "Smallest prefix the provider will actually cache; absent means the provider default applies.", + }, + "token_count_framing": { + **NONNEG_INTEGER, + "description": "Tokens added once per request to a local count estimate, for the provider's chat framing.", + }, + "tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."}, + "rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."}, +} + +NUMBER_KEYS: dict[str, JsonSchema] = { + "regional_processing_uplift_multiplier_eu": { + "type": "number", + "minimum": 1, + "description": "Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%).", + }, + "regional_processing_uplift_multiplier_us": { + "type": "number", + "minimum": 1, + "description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).", + }, + "regional_endpoint_uplift_multiplier": { + "type": "number", + "minimum": 1, + "description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%).", + }, +} + +COST_DESCRIPTIONS: dict[str, str] = { + "input_cost_per_token": "USD per prompt token.", + "output_cost_per_token": "USD per generated token.", + "output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.", + "google_maps_grounding_cost_per_query": ( + "USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit." + ), + "cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.", + "cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.", + "input_cost_per_token_batches": "USD per prompt token via the provider's batch API.", + "output_cost_per_token_batches": "USD per generated token via the provider's batch API.", +} + + +def cost_description(key: str) -> Optional[str]: + if key in COST_DESCRIPTIONS: + return COST_DESCRIPTIONS[key] + if key.endswith("_flex"): + return "Flex service-tier rate for the same-named base field." + if key.endswith("_priority"): + return "Priority service-tier rate for the same-named base field." + if "_above_" in key: + return "Rate applied once the prompt exceeds the token threshold in the field name." + return None + + +def cost_schema(key: str) -> JsonSchema: + description = cost_description(key) + return {**NONNEG_NUMBER, "description": description} if description else dict(NONNEG_NUMBER) + + +def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]: + return { + "litellm_provider": { + "type": "string", + "description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.", + }, + "mode": { + "type": "string", + "description": "Primary API surface / task type of the model.", + "enum": list(modes), + }, + "source": { + "type": "string", + "description": "URL of the provider pricing/model page this entry was taken from.", + }, + "deprecation_date": { + "type": "string", + "description": "Date the provider deprecates the model, YYYY-MM-DD.", + "format": "date", + "pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$", + }, + "web_search_billing_unit": { + "type": "string", + "description": "Whether web search is billed per query or per prompt.", + "enum": ["per_query", "per_prompt"], + }, + "bedrock_output_config_effort_ceiling": { + "type": "string", + "description": "Highest reasoning effort the Bedrock output_config accepts for this model.", + "enum": ["low", "medium", "high", "max", "xhigh"], + }, + "default_reasoning_effort": { + "type": "string", + "description": ( + "Reasoning effort the provider applies when the request omits reasoning_effort. " + "Gates whether a non-default temperature or the top_p/logprobs sampling params are " + "accepted, which hold only when the effort resolves to 'none'." + ), + "enum": ["none", "minimal", "low", "medium", "high", "xhigh"], + }, + "comment": STRING, + "audio_transcription_config": STRING, + "vertex_ai_audio_api": { + "type": "string", + "enum": ["lyria_predict", "lyria_interactions"], + }, + } + + +def classify(key: str, modes: tuple) -> Optional[JsonSchema]: + curated = {**OBJECT_KEYS, **ARRAY_KEYS, **string_key_schemas(modes), **INTEGER_KEYS, **NUMBER_KEYS} + if key in curated: + return curated[key] + if key.startswith("supports_") or key in EXTRA_BOOLEAN_KEYS: + return BOOLEAN + if "cost" in key: + return cost_schema(key) + return None + + +def build_schema(prices: dict) -> JsonSchema: + entries = {name: entry for name, entry in prices.items() if name not in SPECIAL_ROOT_KEYS} + all_keys = tuple(sorted({key for entry in entries.values() for key in entry})) + modes = tuple(sorted({entry["mode"] for entry in entries.values() if "mode" in entry})) + unclassified = tuple(key for key in all_keys if classify(key, modes) is None) + if unclassified: + raise SystemExit( + f"Unclassified keys in {PRICES_PATH.name}: {', '.join(unclassified)}. " + f"Add them to the key tables in {Path(__file__).name} and rerun it." + ) + entry_properties = {key: classify(key, modes) for key in all_keys} + return { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "title": "LiteLLM model_prices_and_context_window.json", + "description": ( + "Schema for LiteLLM's model price and context window registry " + "(https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). " + "Every top-level key except 'sample_spec' and 'fallback_generalizations' is a model id, " + "optionally prefixed with its provider (e.g. 'azure/gpt-5.4'), mapping to a model entry. " + "All costs are USD per unit. New optional fields are added regularly, so consumers should " + "ignore unknown fields rather than reject them." + ), + "type": "object", + "properties": { + "sample_spec": { + "type": "object", + "description": ( + "Documentation placeholder illustrating the entry shape; not a real model and not " + "schema-conformant (several values are prose)." + ), + }, + "fallback_generalizations": { + "type": "object", + "description": "Regex rules that generalize unknown model ids to known families; not a model entry.", + "properties": { + "rules": { + "type": "array", + "items": { + "type": "object", + "properties": { + "name": STRING, + "pattern": STRING, + "description": STRING, + }, + "required": ["name", "pattern"], + "additionalProperties": True, + }, + } + }, + "additionalProperties": False, + }, + }, + "additionalProperties": {"$ref": "#/$defs/modelEntry"}, + "$defs": { + "modelEntry": { + "type": "object", + "description": ( + "Pricing, limits, and capability flags for one model. Fields other than litellm_provider " + "are optional; boolean capability flags are simply omitted when unknown or false." + ), + "required": ["litellm_provider"], + "properties": entry_properties, + "additionalProperties": True, + } + }, + } + + +def render(schema: JsonSchema) -> str: + return json.dumps(schema, indent=2) + "\n" + + +def validation_errors(prices: dict, schema: JsonSchema) -> tuple: + validator = jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER) + return tuple( + f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}" + for error in validator.iter_errors(prices) + ) + + +def main() -> int: + check = "--check" in sys.argv[1:] + prices = json.loads(PRICES_PATH.read_text()) + rendered = render(build_schema(prices)) + errors = validation_errors(prices, json.loads(rendered)) + if errors: + print(f"{PRICES_PATH.name} does not validate against the generated schema:") + print("\n".join(errors[:20])) + return 1 + if not check: + SCHEMA_PATH.write_text(rendered) + print(f"wrote {SCHEMA_PATH}") + return 0 + if not SCHEMA_PATH.exists() or SCHEMA_PATH.read_text() != rendered: + print( + f"{SCHEMA_PATH.name} is out of sync with {PRICES_PATH.name}. " + f"Run `python {Path(__file__).relative_to(REPO_ROOT)}` and commit the result." + ) + return 1 + print(f"{SCHEMA_PATH.name} is in sync and {PRICES_PATH.name} validates against it") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/litellm-rust/crates/model-catalog/AGENTS.md b/litellm-rust/crates/model-catalog/AGENTS.md index 0a380788c19..9fcbcd57a76 100644 --- a/litellm-rust/crates/model-catalog/AGENTS.md +++ b/litellm-rust/crates/model-catalog/AGENTS.md @@ -2,4 +2,5 @@ For `model_prices_and_context_window.json` validation, we should eventually: +- Remove any schema file like `model_prices_and_context_window.schema.json` - Stop skipping this crate's tests diff --git a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py deleted file mode 100644 index 3c437ca7fbb..00000000000 --- a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py +++ /dev/null @@ -1,37 +0,0 @@ -import json -import re -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass -from typing import Final - - -@dataclass(frozen=True, slots=True) -class ChatTemplateToolSection: - preamble_tokens: int - implicit_system_turn_tokens: int - render_tool_schemas: Callable[[Sequence[Mapping[str, object]]], str] - - -def _deepseek_v4_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: - return "\n".join(json.dumps(function, ensure_ascii=False) for function in functions) - - -def _qwen3_5_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: - return "".join( - "\n" + json.dumps({"type": "function", "function": function}, ensure_ascii=False) for function in functions - ) - - -_TOOL_SECTIONS: Final = ( - ( - re.compile(r"deepseek-v4(?![.\d]|p\d)", re.IGNORECASE), - ChatTemplateToolSection(222, 0, _deepseek_v4_tool_schemas), - ), - (re.compile(r"qwen3[._]?5(?!\d)", re.IGNORECASE), ChatTemplateToolSection(196, 5, _qwen3_5_tool_schemas)), -) - - -def chat_template_tool_section(model: str | None) -> ChatTemplateToolSection | None: - if not model: - return None - return next((section for pattern, section in _TOOL_SECTIONS if pattern.search(model)), None) diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py index 37c32144bbf..b5e9afca86b 100644 --- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py +++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py @@ -31,7 +31,7 @@ Anthropic wire shape is built later by ``anthropic_messages_pt``. from collections.abc import Iterator, Mapping, Sequence from itertools import chain, groupby -from typing import Final, Literal, TypeAlias, TypeVar +from typing import Final, Literal, TypeAlias from litellm.types.llms.anthropic import AnthropicMessagesSystemMessageParam, AnthropicSystemMessageContent from litellm.types.llms.openai import ( @@ -55,7 +55,6 @@ _RENDERED_ASSISTANT_PART_TYPES: Final = frozenset({"text", "server_tool_use"}) _THINKING_BLOCK_TYPES: Final = frozenset({"thinking", "redacted_thinking"}) _MessageKind: TypeAlias = Literal["system", "tool", "user", "other"] -_Message: Final = TypeVar("_Message") _TextPart: TypeAlias = tuple[str, ChatCompletionCachedContent | None] @@ -97,8 +96,8 @@ def _kind(message: object) -> _MessageKind: def split_leading_system_run( - messages: Sequence[_Message], -) -> tuple[tuple[_Message, ...], tuple[_Message, ...]]: + messages: Sequence[AllMessageValues], +) -> tuple[tuple[AllMessageValues, ...], tuple[AllMessageValues, ...]]: """Split ``messages`` into the leading run of system messages and everything after it.""" leading_count: Final = next( (index for index, message in enumerate(messages) if not is_system_message(message)), @@ -161,16 +160,9 @@ def _anthropic_text_block(part: _TextPart) -> AnthropicSystemMessageContent: return cached -def anthropic_system_blocks(run: Sequence[object]) -> tuple[AnthropicSystemMessageContent, ...]: - """The top-level ``system`` blocks for a run of system messages: every non-empty text part in order, - each keeping its ``cache_control``, which is the shape the chat path sends for the leading run.""" - parts: Final = chain.from_iterable(_text_parts(message) for message in run) - return tuple(_anthropic_text_block(part) for part in parts) - - def anthropic_system_messages(message: object) -> tuple[AnthropicMessagesSystemMessageParam, ...]: """The Anthropic wire message for a system message, or nothing when it carries no text.""" - blocks: Final = anthropic_system_blocks((message,)) + blocks: Final = tuple(_anthropic_text_block(part) for part in _text_parts(message)) if not blocks: return () wire: Final[AnthropicMessagesSystemMessageParam] = { diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 5a76e85da75..c62233f49e2 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -34,10 +34,6 @@ from litellm.constants import ( TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.prompt_templates.chat_template_tools import ( - ChatTemplateToolSection, - chat_template_tool_section, -) from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace, HuggingFaceTokenizer, OpenAIEncoding from litellm.litellm_core_utils.url_utils import safe_get from litellm.llms.custom_httpx.http_handler import _get_httpx_client @@ -347,9 +343,6 @@ EXTRAPOLATION_SAMPLES: Final = 16 T_ParamSpec: Final = ParamSpec("T_ParamSpec") T_Retval = TypeVar("T_Retval") _COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter") -_ANTHROPIC_TOOL_KEYS_MAPPED_BY_BRIDGE: Final = frozenset( - ("name", "input_schema", "description", "cache_control", "strict", "type", "eager_input_streaming") -) def _count_offload_limiter_for_this_loop() -> anyio.CapacityLimiter: @@ -407,7 +400,6 @@ class _MessageCountParams: self.tokens_per_message = 3 self.tokens_per_name = 1 self.count_function = _get_count_function(model, custom_tokenizer) - self.tool_section = chat_template_tool_section(model) def token_counter( @@ -471,9 +463,7 @@ def token_counter( num_tokens = _count_messages(params, new_messages, use_default_image_token_count, default_token_count) if count_response_tokens is False: includes_system_message: Final = any(message.get("role", None) == "system" for message in new_messages) - num_tokens += _count_extra( - params.count_function, tools, tool_choice, includes_system_message, params.tool_section - ) + num_tokens += _count_extra(params.count_function, tools, tool_choice, includes_system_message) else: raise ValueError("Either text or messages must be provided") @@ -567,7 +557,6 @@ def _count_extra( tools: list[ChatCompletionToolParam] | None, tool_choice: ChatCompletionNamedToolChoiceParam | None, includes_system_message: bool, - tool_section: ChatTemplateToolSection | None, ) -> int: """Count extra tokens for function definitions and tool choices. Args: @@ -575,56 +564,26 @@ def _count_extra( tools (Optional[List[ChatCompletionToolParam]]): The available tools. tool_choice (Optional[ChatCompletionNamedToolChoiceParam]): The tool choice. includes_system_message (bool): Whether the messages include a system message. - tool_section (Optional[ChatTemplateToolSection]): How the model's chat template renders tools, if known. """ - reply_primer_tokens: Final = 3 - tool_choice_tokens: Final = 0 if tool_section is not None else _count_tool_choice(count_function, tool_choice) - return ( - reply_primer_tokens - + _count_tool_definitions(count_function, tools, includes_system_message, tool_section) - + tool_choice_tokens - ) + num_tokens = 3 # every reply is primed with <|start|>assistant<|message|> -def _count_tool_definitions( - count_function: TokenCounterFunction, - tools: Sequence[object] | None, - includes_system_message: bool, - tool_section: ChatTemplateToolSection | None, -) -> int: - if not tools: - return 0 - if tool_section is not None: - return _count_chat_template_tools(count_function, tools, includes_system_message, tool_section) - system_message_discount: Final = 4 if includes_system_message else 0 - return count_function(_format_function_definitions(tools)) + 9 - system_message_discount - - -def _count_tool_choice( - count_function: TokenCounterFunction, - tool_choice: ChatCompletionNamedToolChoiceParam | None, -) -> int: + if tools: + num_tokens += count_function(_format_function_definitions(tools)) + num_tokens += 9 # Additional tokens for function definition of tools + # If there's a system message and tools are present, subtract four tokens + if tools and includes_system_message: + num_tokens -= 4 + # If tool_choice is 'none', add one token. + # If it's an object, add 4 + the number of tokens in the function name. + # If it's undefined or 'auto', don't add anything. if tool_choice == "none": - return 1 - if isinstance(tool_choice, dict): - return 7 + count_function(str(tool_choice["function"]["name"])) - return 0 + num_tokens += 1 + elif isinstance(tool_choice, dict): + num_tokens += 7 + num_tokens += count_function(str(tool_choice["function"]["name"])) - -def _count_chat_template_tools( - count_function: TokenCounterFunction, - tools: Sequence[object], - includes_system_message: bool, - tool_section: ChatTemplateToolSection, -) -> int: - functions: Final = tuple( - function for function in _function_definitions(_tools_sent_as_functions(tools)) if function.get("name") - ) - if not functions: - return 0 - system_turn_tokens: Final = 0 if includes_system_message else tool_section.implicit_system_turn_tokens - schemas: Final = tool_section.render_tool_schemas(functions) - return tool_section.preamble_tokens + system_turn_tokens + count_function(schemas) + return num_tokens def _get_extrapolating_count_function( @@ -1034,90 +993,49 @@ def _count_content_list( ) -def _format_function_definitions(tools: Sequence[object]) -> str: +def _format_function_definitions(tools): """Formats tool definitions in the format that OpenAI appears to use. Based on https://github.com/forestwanglin/openai-java/blob/main/jtokkit/src/main/java/xyz/felh/openai/jtokkit/utils/TikTokenUtils.java """ lines: Final = [] lines.append("namespace functions {") lines.append("") - for function in _function_definitions(tools): - lines.extend(_format_single_function_definition(function)) + for tool in tools: + if not isinstance(tool, dict): + continue + function = tool.get("function") + if not isinstance(function, dict): + # Anthropic tool shape → OpenAI function dict for token counting. + params = tool.get("input_schema") or tool.get("parameters") or {} + if not isinstance(params, dict): + params = {} + function = { + "name": tool.get("name"), + "description": tool.get("description"), + "parameters": params, + } + function_name = function.get("name") + if not function_name: + # Skip malformed tools missing a name to avoid emitting + # ``type None = ...`` which would produce inaccurate token counts. + continue + if function_description := function.get("description"): + lines.append(f"// {function_description}") + parameters = function.get("parameters") or {} + if not isinstance(parameters, dict): + parameters = {} + properties = parameters.get("properties") + if properties and properties.keys(): + lines.append(f"type {function_name} = (_: {{") + lines.append(_format_object_parameters(parameters, 0)) + lines.append("}) => any;") + else: + lines.append(f"type {function_name} = () => any;") + lines.append("") lines.append("} // namespace functions") return "\n".join(lines) -def _tools_sent_as_functions(tools: Sequence[object]) -> tuple[Mapping[str, object], ...]: - return tuple(tool for tool in _tool_mappings(tools) if not _is_anthropic_web_search_tool(tool)) - - -def _is_anthropic_web_search_tool(tool: Mapping[str, object]) -> bool: - return tool.get("name") == "web_search" and "input_schema" not in tool and tool.get("type") != "function" - - -def _tool_mappings(tools: Sequence[object]) -> Iterable[Mapping[str, object]]: - for tool in tools: - if isinstance(tool, Mapping): - yield cast(Mapping[str, object], tool) - - -def _function_definitions(tools: Sequence[object]) -> Iterable[Mapping[str, object]]: - for tool in _tool_mappings(tools): - yield from _function_definitions_for_tool(tool) - - -def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mapping[str, object]]: - function: Final = tool.get("function") - if isinstance(function, Mapping): - yield function - return - declarations: Final = tool.get("function_declarations") or tool.get("functionDeclarations") - if isinstance(declarations, list): - for declaration in declarations: - if isinstance(declaration, Mapping): - yield declaration - return - input_schema: Final = tool.get("input_schema") - parameters: Final = input_schema or tool.get("parameters") or {} - normalized_parameters: Final = parameters if isinstance(parameters, Mapping) else {} - unmapped_anthropic_keys: Final = ( - {key: value for key, value in tool.items() if key not in _ANTHROPIC_TOOL_KEYS_MAPPED_BY_BRIDGE} - if input_schema is not None - else {} - ) - description: Final = tool.get("description") - strict: Final = tool.get("strict") - yield { - "name": tool.get("name"), - "parameters": {**normalized_parameters, **unmapped_anthropic_keys}, - **({"description": description} if description is not None else {}), - **({"strict": bool(strict)} if strict is not None else {}), - } - - -def _format_single_function_definition(function: Mapping[str, object]) -> tuple[str, ...]: - function_name: Final = function.get("name") - if not function_name: - return () - function_description: Final = function.get("description") - parameters_value: Final = function.get("parameters") or {} - parameters: Final = parameters_value if isinstance(parameters_value, Mapping) else {} - properties: Final = parameters.get("properties") - if isinstance(properties, Mapping) and properties: - return ( - *((f"// {function_description}",) if function_description else ()), - f"type {function_name} = (_: {{", - _format_object_parameters(parameters, 0), - "}) => any;", - "", - ) - return ( - *((f"// {function_description}",) if function_description else ()), - f"type {function_name} = () => any;", - "", - ) - - def _format_object_parameters(parameters, indent): properties: Final = parameters.get("properties") if not properties: diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index 0cecf62368c..dd2135f4918 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -12,7 +12,6 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.asyncify import asyncify from litellm.llms.anthropic.common_utils import AnthropicError from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, @@ -63,7 +62,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): verbose_logger.debug("Processing Anthropic CountTokens request for model: %s", model) # Transform request to Anthropic format - request_body: Final = await asyncify(self.transform_request_to_count_tokens)( + request_body: Final = self.transform_request_to_count_tokens( model=model, messages=messages, tools=tools, diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 8c47d1cf884..fb12747cec0 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -11,67 +11,11 @@ from typing import Final from pydantic import JsonValue, TypeAdapter from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION -from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( - anthropic_system_blocks, - split_leading_system_run, -) -from litellm.types.llms.openai import ChatCompletionImageObject _COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue]) -_SYSTEM_BLOCKS: Final = TypeAdapter(list[JsonValue]) -_IMAGE_BLOCK: Final = TypeAdapter(ChatCompletionImageObject) COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config") -def _count_image(block: JsonValue) -> JsonValue: - if not isinstance(block, dict) or block.get("type") != "image_url": - return block - from litellm.litellm_core_utils.prompt_templates.factory import convert_to_anthropic_image_obj - - image_block: Final = _IMAGE_BLOCK.validate_python(block) - image_url: Final = image_block["image_url"] - source: Final = convert_to_anthropic_image_obj( - openai_image_url=image_url if isinstance(image_url, str) else image_url["url"], - format=image_url.get("format") if isinstance(image_url, dict) else None, - ) - image: Final = _COUNT_REQUEST.validate_python({"type": "image", "source": source}) - return {**{key: value for key, value in block.items() if key not in {"type", "image_url"}}, **image} - - -def _count_block(block: JsonValue) -> JsonValue: - if not isinstance(block, dict) or block.get("type") != "tool_result": - return _count_image(block) - content: Final = block.get("content") - if not isinstance(content, list): - return block - return {**block, "content": [_count_image(part) for part in content]} - - -def _count_content(content: JsonValue) -> JsonValue: - return [_count_block(block) for block in content] if isinstance(content, list) else content - - -def _lift_leading_system( - messages: Sequence[Mapping[str, JsonValue]], system: JsonValue -) -> tuple[tuple[Mapping[str, JsonValue], ...], JsonValue]: - """Move the leading run of system-role messages into the top-level ``system`` parameter. - - count_tokens only takes the initial system prompt there and answers 400 on ``role: "system"`` - at the head of ``messages``; the chat path sends the same run as ``system``. A caller's own - ``system`` keeps its place ahead of the lifted blocks, and a ``system`` that is neither text - nor a block list is left as sent, messages included, for the provider to judge. - """ - leading, conversation = split_leading_system_run(messages) - if not leading or not (system is None or isinstance(system, (str, list))): - return tuple(messages), system - lifted: Final = _SYSTEM_BLOCKS.validate_python(list(anthropic_system_blocks(leading))) - if isinstance(system, list): - return conversation, [*system, *lifted] - if isinstance(system, str) and system: - return conversation, [{"type": "text", "text": system}, *lifted] - return conversation, lifted or system - - class AnthropicCountTokensConfig: """ Configuration and transformation logic for Anthropic CountTokens API. @@ -102,24 +46,16 @@ class AnthropicCountTokensConfig: """ Transform request to Anthropic CountTokens format. - Includes optional system and tools fields for accurate token counting; a leading run of - system-role messages is counted through ``system``, the only place count_tokens accepts it. + Includes optional system and tools fields for accurate token counting. """ options: Final[Mapping[str, JsonValue]] = optional_params or MappingProxyType({}) - counted_messages, counted_system = _lift_leading_system(messages, system) return _COUNT_REQUEST.validate_python( MappingProxyType( { "model": model, - "messages": [ - {**message, "content": _count_content(message["content"])} for message in counted_messages - ], + "messages": messages, **MappingProxyType( - { - key: value - for key, value in (("system", counted_system), ("tools", tools)) - if value is not None - } + {key: value for key, value in (("system", system), ("tools", tools)) if value is not None} ), **MappingProxyType( {key: value for key, value in options.items() if key in COUNT_TOKEN_OPTION_NAMES} diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index e1eb83e4ebf..36a227c03ee 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -10,7 +10,6 @@ import httpx import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.asyncify import asyncify from litellm.llms.anthropic.common_utils import AnthropicError from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( AzureAIAnthropicCountTokensConfig, @@ -60,7 +59,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): verbose_logger.debug("Processing Azure AI Anthropic CountTokens request for model: %s", model) # Transform request to Anthropic format - request_body: Final = await asyncify(self.transform_request_to_count_tokens)( + request_body: Final = self.transform_request_to_count_tokens( model=model, messages=messages, tools=tools, diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 7ff2fddbb57..57bbb5b6cb6 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -3,101 +3,29 @@ Bedrock Token Counter implementation using the CountTokens API. """ from collections.abc import Mapping, Sequence -from dataclasses import dataclass from typing import Any, Final from pydantic import JsonValue -from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler -from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler -from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse - -RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" -MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" -_RUNTIME_FALLBACK_STATUSES: Final = frozenset({400, 404}) - - -class _CountTokensReply(LiteLLMPydanticObjectBase): - input_tokens: int - - -@dataclass(frozen=True, slots=True) -class CountedTokens: - input_tokens: int - original_response: Mapping[str, JsonValue] - tokenizer_type: str - - -@dataclass(frozen=True, slots=True) -class CountTokensFailure: - status_code: int - message: str - tokenizer_type: str - - -CountTokensOutcome = CountedTokens | CountTokensFailure - - -def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool: - return ( - isinstance(outcome, CountTokensFailure) - and outcome.status_code in _RUNTIME_FALLBACK_STATUSES - and is_mantle_claude_model(resolved_model) - ) +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse class BedrockTokenCounter(BaseTokenCounter): - """Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens - on bedrock-mantle for Claude models bedrock-runtime answers 400 or 404 for""" - - def __init__( - self, - runtime_handler: BedrockCountTokensHandler | None = None, - mantle_handler: BedrockMantleCountTokensHandler | None = None, - client: AsyncHTTPHandler | None = None, - ) -> None: - self._runtime_handler: Final = runtime_handler or BedrockCountTokensHandler() - self._mantle_handler: Final = mantle_handler or BedrockMantleCountTokensHandler() - self._client: Final = client + """Token counter implementation for AWS Bedrock provider using the CountTokens API.""" def should_use_token_counting_api( self, custom_llm_provider: str | None = None, ) -> bool: + """ + Returns True if we should use the Bedrock CountTokens API for token counting. + """ return custom_llm_provider == LlmProviders.BEDROCK.value - async def _count_with( - self, - handler: BedrockCountTokensHandler | BedrockMantleCountTokensHandler, - tokenizer_type: str, - request_data: dict[str, object], - litellm_params: dict[str, object], - resolved_model: str, - ) -> CountTokensOutcome: - try: - result: Final = await handler.handle_count_tokens_request( - request_data=request_data, - litellm_params=litellm_params, - resolved_model=resolved_model, - client=self._client, - ) - reply: Final = _CountTokensReply.model_validate(result) - except BedrockError as e: - verbose_logger.debug( - "%s CountTokens API error: status=%s, message=%s", tokenizer_type, e.status_code, e.message - ) - return CountTokensFailure(status_code=e.status_code, message=e.message, tokenizer_type=tokenizer_type) - except Exception as e: - verbose_logger.debug("Error calling %s CountTokens API: %s", tokenizer_type, e) - return CountTokensFailure(status_code=500, message=str(e), tokenizer_type=tokenizer_type) - return CountedTokens(input_tokens=reply.input_tokens, original_response=result, tokenizer_type=tokenizer_type) - async def count_tokens( self, model_to_use: str, @@ -110,52 +38,81 @@ class BedrockTokenCounter(BaseTokenCounter): messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, ) -> TokenCountResponse | None: + """ + Count tokens using AWS Bedrock's CountTokens API. + + This method calls the existing BedrockCountTokensHandler to make an API call + to Bedrock's token counting endpoint, bypassing the local tiktoken-based counting. + + Args: + model_to_use: The model identifier + messages: The messages to count tokens for + contents: Alternative content format (not used for Bedrock) + deployment: Deployment configuration containing litellm_params + request_model: The original request model name + + Returns: + TokenCountResponse with token count, or None if counting fails + """ if not messages: return None - litellm_params: Final = (deployment or {}).get("litellm_params", {}) + deployment = deployment or {} + litellm_params: Final = deployment.get("litellm_params", {}) + + # Build request data in the format expected by BedrockCountTokensHandler request_data: Final[dict[str, object]] = { "model": model_to_use, "messages": messages, - **({"tools": tools} if tools else {}), - **({"system": system} if system else {}), } + + if tools: + request_data["tools"] = tools + + if system: + request_data["system"] = system + + # Get the resolved model (strip prefixes like bedrock/, converse/, etc.) resolved_model: Final = get_bedrock_base_model(model_to_use) - runtime: Final = await self._count_with( - self._runtime_handler, RUNTIME_TOKENIZER_TYPE, request_data, litellm_params, resolved_model - ) - outcome: Final = ( - await self._count_with( - self._mantle_handler, MANTLE_TOKENIZER_TYPE, request_data, litellm_params, resolved_model + try: + handler: Final = BedrockCountTokensHandler() + result: Final = await handler.handle_count_tokens_request( + request_data=request_data, + litellm_params=litellm_params, + resolved_model=resolved_model, ) - if _runtime_rejected_claude_model(runtime, resolved_model) - else runtime - ) - match outcome: - case CountedTokens(): + + # Transform response to TokenCountResponse + if result is not None: return TokenCountResponse( - total_tokens=outcome.input_tokens, + total_tokens=result.get("input_tokens", 0), request_model=request_model, model_used=model_to_use, - tokenizer_type=outcome.tokenizer_type, - original_response=dict(outcome.original_response), + tokenizer_type="bedrock_api", + original_response=result, ) - case CountTokensFailure(): - verbose_logger.warning( - "%s CountTokens API error: status=%s, message=%s", - outcome.tokenizer_type, - outcome.status_code, - outcome.message, - ) - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type=outcome.tokenizer_type, - error=True, - error_message=outcome.message, - status_code=outcome.status_code, - ) - case _: - assert_never(outcome) + except BedrockError as e: + verbose_logger.warning("Bedrock CountTokens API error: status=%s, message=%s", e.status_code, e.message) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="bedrock_api", + error=True, + error_message=e.message, + status_code=e.status_code, + ) + except Exception as e: + verbose_logger.warning("Error calling Bedrock CountTokens API: %s", e) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="bedrock_api", + error=True, + error_message=str(e), + status_code=500, + ) + + return None diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index be8243e48ef..d7fb510f057 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -125,7 +125,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): raise except httpx.HTTPStatusError as e: # HTTP errors - preserve the actual status code - verbose_logger.debug("HTTP error in CountTokens handler: %s", e) + verbose_logger.error("HTTP error in CountTokens handler: %s", e) raise BedrockError( status_code=e.response.status_code, message=e.response.text, diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py deleted file mode 100644 index 23c3dc54c95..00000000000 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ /dev/null @@ -1,99 +0,0 @@ -from typing import Final - -import httpx -from pydantic import JsonValue, TypeAdapter -from typing_extensions import NotRequired, ReadOnly, TypedDict - -import litellm -from litellm._logging import verbose_logger -from litellm.llms.anthropic.count_tokens.transformation import AnthropicCountTokensConfig -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing -from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client - -MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens" -MANTLE_ANTHROPIC_VERSION: Final = "2023-06-01" - - -class MantleCountTokensRequest(TypedDict): - messages: ReadOnly[list[dict[str, JsonValue]]] - system: ReadOnly[NotRequired[JsonValue]] - tools: ReadOnly[NotRequired[list[dict[str, JsonValue]]]] - - -_COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest) -_COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue]) - - -class BedrockMantleCountTokensHandler(BaseAWSLLM): - """Counts tokens through Anthropic's count_tokens on the bedrock-mantle endpoint. - - Claude models that Bedrock offers only through cross-region inference answer 400 on - bedrock-runtime's CountTokens; AWS documents Mantle's /anthropic/v1/messages/count_tokens - as the way to count them, with the base model id and the deployment's AWS credentials - """ - - def __init__(self, anthropic_config: AnthropicCountTokensConfig | None = None) -> None: - super().__init__() - self._anthropic_config: Final = anthropic_config or AnthropicCountTokensConfig() - - def get_mantle_count_tokens_endpoint(self, aws_region_name: str) -> str: - messages_url: Final = build_mantle_messages_url( - api_base=None, aws_bedrock_runtime_endpoint=None, region=aws_region_name - ) - return f"{messages_url}{MANTLE_COUNT_TOKENS_SUFFIX}" - - async def handle_count_tokens_request( - self, - request_data: dict[str, object], - litellm_params: dict[str, object], - resolved_model: str, - client: AsyncHTTPHandler | None = None, - ) -> dict[str, JsonValue]: - try: - request: Final = _COUNT_REQUEST.validate_python(request_data) - aws_region_name: Final = self._get_aws_region_name( - optional_params=litellm_params, model=resolved_model, model_id=None - ) - endpoint_url: Final = self.get_mantle_count_tokens_endpoint(aws_region_name) - body: Final = self._anthropic_config.transform_request_to_count_tokens( - model=resolved_model, - messages=request["messages"], - tools=request.get("tools"), - system=request.get("system"), - ) - verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url) - api_key: Final = litellm_params.get("api_key") - signed_headers, signed_body = await run_aws_signing( - self._sign_request, - service_name="bedrock", - headers={"Content-Type": "application/json", "anthropic-version": MANTLE_ANTHROPIC_VERSION}, - optional_params=litellm_params, - request_data=body, - api_base=endpoint_url, - model=resolved_model, - api_key=api_key if isinstance(api_key, str) else None, - ) - async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) - response: Final = await async_client.post( - endpoint_url, headers=signed_headers, data=signed_body, timeout=30.0 - ) - if response.status_code != 200: - raise BedrockError( - status_code=response.status_code, - message=response.text, - headers=response.headers, - response=response, - ) - return _COUNT_RESPONSE.validate_json(response.content) - except BedrockError: - raise - except httpx.HTTPStatusError as e: - raise BedrockError( - status_code=e.response.status_code, - message=e.response.text, - headers=e.response.headers, - response=e.response, - ) - except Exception as e: - raise BedrockError(status_code=500, message=f"bedrock-mantle count_tokens processing error: {e}") diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index ca637af0110..9892ef6224e 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -158,7 +158,3 @@ def mantle_base_segment(model: str | None, model_cost: dict) -> str: """ entry: Final = model_cost.get(f"bedrock_mantle/{split_mantle_region_prefix(model)[1]}", {}) if model else {} return "openai/v1" if entry.get("use_openai_responses_path") is True else "v1" - - -def is_mantle_claude_model(model: str) -> bool: - return "claude" in model.lower() diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 3ea5cf5c4b7..a26aebd8526 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -193,10 +193,9 @@ class OpenAICountTokensConfig: tools: list[dict[str, Any]], ) -> list[dict[str, object]]: """ - Transform OpenAI chat and Anthropic tools to Responses API tools format. + Transform OpenAI chat tools format to Responses API tools format. Chat format: {"type": "function", "function": {"name": "...", "parameters": {...}}} - Anthropic format: {"name": "...", "input_schema": {...}} Responses format: {"type": "function", "name": "...", "parameters": {...}} """ transformed: Final[list[dict[str, object]]] = [] @@ -212,16 +211,6 @@ class OpenAICountTokensConfig: if "strict" in func: item["strict"] = func["strict"] transformed.append(item) - elif "input_schema" in tool and "name" in tool and tool.get("type") in (None, "custom"): - transformed.append( - { - "type": "function", - "name": tool["name"], - "description": tool.get("description", ""), - "parameters": tool["input_schema"], - "strict": bool(tool.get("strict")), - } - ) else: # Pass through non-function tools (e.g., web_search, file_search) transformed.append(tool) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index eb57ef4e6c6..2c7e07616b5 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1242,7 +1242,6 @@ class VertexAITokenCounter(BaseTokenCounter): ) -> TokenCountResponse | None: import copy - from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( VertexAIError as PartnerVertexAIError, ) @@ -1253,23 +1252,6 @@ class VertexAITokenCounter(BaseTokenCounter): deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) - without_count_route: Final = ( - VertexAIPartnerModels.is_vertex_partner_model(model_to_use) and count_tokens_publisher(model_to_use) is None - ) - model_info: Final[Mapping[str, object]] = deployment.get("model_info") or {} - if without_count_route and messages and model_info.get("custom_tokenizer") is not None: - return None - if without_count_route: - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="vertex_ai_partner_models", - error=True, - error_message=f"LiteLLM does not support counting tokens for Vertex AI model {model_to_use}.", - status_code=400, - ) - # Check if this is a partner model (Claude, Mistral, etc.) if VertexAIPartnerModels.is_vertex_partner_model(model_to_use): # Use partner models token counter diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 7100c9efde3..285380f5fc7 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -1,5 +1,5 @@ """ -Token counter for Vertex AI Partner Models (Anthropic Claude, Meta Llama) +Token counter for Vertex AI Partner Models (Anthropic Claude, Mistral, etc.) This handler provides token counting for partner models hosted on Vertex AI. Unlike Gemini models which use Google's token counting API, partner models use @@ -12,30 +12,36 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -_COUNT_TOKENS_PUBLISHERS: Final = ( - ("claude", "anthropic"), - ("llama", "meta"), - ("meta/", "meta"), -) - - -def count_tokens_publisher(model: str) -> str | None: - return next((publisher for marker, publisher in _COUNT_TOKENS_PUBLISHERS if marker in model), None) - class VertexAIPartnerModelsTokenCounter(VertexBase): """ Token counter for Vertex AI Partner Models. - Handles token counting for models like Claude (Anthropic) and Llama (Meta) + Handles token counting for models like Claude (Anthropic), Mistral, etc. that are available through Vertex AI's partner model program. """ def _get_publisher_for_model(self, model: str) -> str: - publisher: Final = count_tokens_publisher(model) - if publisher is None: + """ + Determine the publisher name for the given model. + + Args: + model: The model name (e.g., "claude-3-5-sonnet-20241022") + + Returns: + Publisher name to use in the Vertex AI endpoint URL + + Raises: + ValueError: If the model is not a recognized partner model + """ + if "claude" in model: + return "anthropic" + elif "mistral" in model or "codestral" in model: + return "mistralai" + elif "llama" in model or "meta/" in model: + return "meta" + else: raise ValueError(f"Unknown partner model: {model}") - return publisher def _build_count_tokens_endpoint( self, diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index b24761ac039..202c61b620e 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -42,10 +42,6 @@ def openai_error_type(exc: object, status_code: int) -> str: carried: Final = attribute_of(exc, "type") if isinstance(carried, str) and carried != STRINGIFIED_NONE: return carried - return openai_error_type_for_status(status_code) - - -def openai_error_type_for_status(status_code: int) -> str: mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) if mapped is not None: return mapped diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8369e79596c..4ec1cf36408 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14153,10 +14153,8 @@ async def count_request_tokens( model_info.get("custom_tokenizer", None), ) system_message: Final = _system_message(system) - prompt_turn: Final = 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 messages), + Sequence[AllMessageValues] | None, messages ) counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index d33d8a83bfa..80652e63300 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -14,12 +14,11 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse from openai.types.responses import ResponseItemList from openai.types.responses.response_create_params import ResponseInputParam -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, with_config +from pydantic import BaseModel, ConfigDict, ValidationError from starlette.websockets import WebSocket, WebSocketDisconnect -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping from litellm.constants import EMPTY_MAPPING from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -32,16 +31,10 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth_websocket, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing -from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_set_request_parsed_body, ) -from litellm.proxy.common_utils.openai_error_payload import ( - LITELLM_CALL_ID_HEADER, - error_status_code, - openai_error_type_for_status, -) from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.types.llms.openai import ( @@ -1112,85 +1105,6 @@ def _responses_invalid_request_response(message: str, param: str | None, code: s return JSONResponse(status_code=400, content=body) -class _ProviderErrorMessage(TypedDict): - message: ReadOnly[str] - - -class _NestedProviderError(TypedDict): - error: ReadOnly[_ProviderErrorMessage] - - -@with_config(ConfigDict(extra="ignore")) -class _ProviderOpenAIErrorDetail(TypedDict): - message: ReadOnly[str] - type: ReadOnly[str] - param: NotRequired[ReadOnly[str | None]] - code: NotRequired[ReadOnly[str | None]] - - -@with_config(ConfigDict(extra="forbid")) -class _ProviderOpenAIError(TypedDict): - error: ReadOnly[_ProviderOpenAIErrorDetail] - - -_PROVIDER_COUNT_ERROR_TYPE: Final = "token_counting_error" -_PROVIDER_OPENAI_ERROR: Final = TypeAdapter(_ProviderOpenAIError) -_NESTED_PROVIDER_ERROR: Final = TypeAdapter(_NestedProviderError) - - -def _provider_error_json(raw_message: str) -> str: - start: Final = raw_message.find("{") - return raw_message if start == -1 else raw_message[start:] - - -def _provider_error_message(raw_message: str) -> str: - error_json: Final = _provider_error_json(raw_message) - try: - return _NESTED_PROVIDER_ERROR.validate_json(error_json)["error"]["message"] - except ValidationError: - flat_message: Final = AnthropicExceptionMapping.extract_error_message(error_json) - return raw_message if flat_message == error_json else flat_message - - -def _provider_count_error_body(raw_message: str, status_code: int) -> _ResponsesApiErrorBody: - try: - provider_error: Final = _PROVIDER_OPENAI_ERROR.validate_json(_provider_error_json(raw_message))["error"] - except ValidationError: - return { - "error": { - "message": _provider_error_message(raw_message), - "type": openai_error_type_for_status(status_code), - "param": None, - "code": None, - } - } - return { - "error": { - "message": provider_error["message"], - "type": provider_error["type"], - "param": provider_error.get("param"), - "code": provider_error.get("code"), - } - } - - -def _provider_count_error_response(exc: ProxyException, request: Request) -> JSONResponse: - from litellm.proxy.proxy_server import ( - _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; a route that answers an error itself must stamp the OTel server span like the global ProxyException handler does - general_settings_view, - ) - - status_code: Final = error_status_code(exc, 500) - _close_dangling_otel_server_span(request, status_code, exc=exc) - body: Final = _provider_count_error_body(exc.message, status_code) - call_id: Final = error_body_call_id(general_settings_view(), exc.headers.get(LITELLM_CALL_ID_HEADER)) - return JSONResponse( - status_code=status_code, - content={"error": with_call_id({**body["error"]}, call_id)}, - headers=exc.headers, - ) - - def _missing_responses_param_response(param: str) -> JSONResponse: return _responses_invalid_request_response( message=f"Missing required parameter: '{param}'.", @@ -1282,12 +1196,7 @@ async def responses_input_tokens( message=f"Invalid request for token counting: {e}", param=None, code=None ) - try: - token_response: Final = await token_counter(request=token_request, call_endpoint=True) - except ProxyException as e: - if e.type != _PROVIDER_COUNT_ERROR_TYPE: - raise - return _provider_count_error_response(e, request) + token_response: Final = await token_counter(request=token_request, call_endpoint=True) result: Final[_ResponsesInputTokensResult] = { "object": "response.input_tokens", "input_tokens": token_response.total_tokens, diff --git a/tests/integration/_support/bedrock_runtime_peer.py b/tests/integration/_support/bedrock_runtime_peer.py deleted file mode 100644 index 64115d9eaa0..00000000000 --- a/tests/integration/_support/bedrock_runtime_peer.py +++ /dev/null @@ -1,278 +0,0 @@ -from __future__ import annotations - -import json -import re -import threading -from collections.abc import Mapping -from multiprocessing.sharedctypes import Synchronized -from types import MappingProxyType -from typing import Final -from urllib.parse import unquote - -from integration._support.upstream import _aws_event_frame -from integration._support.wire import Reply, Request, wire_server -from pydantic import JsonValue, TypeAdapter - -MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") -EVENT_STREAM: Final = "application/vnd.amazon.eventstream" -REASONING_EFFORTS: Final = frozenset(("none", "minimal", "low", "medium", "high", "xhigh")) -NATIVE_CHAT: Final = "/openai/v1/chat/completions" -NATIVE_RESPONSES: Final = "/openai/v1/responses" -PNG_1X1: Final = bytes.fromhex( - "89504e470d0a1a0a0000000d49484452000000010000000108060000001f15c489" - "0000000d49444154789c63f8cfc0f01f00050001ff89993d1d0000000049454e44ae426082" -) -USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType( - { - "prompt_tokens": 9, - "completion_tokens": 5, - "total_tokens": 14, - "completion_tokens_details": {"reasoning_tokens": 3}, - } -) -_STATUS: Final = re.compile(r"status=(\d{3})") -_CONVERSE: Final = re.compile(r"^/model/(.+)/converse$") -_CONVERSE_STREAM: Final = re.compile(r"^/model/(.+)/converse-stream$") -_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) -_NO_MARKER: Final = "0" * 32 - - -def marker_of(request: Request) -> str: - found: Final = MARKER.search(request.body.decode(errors="replace")) - return _NO_MARKER if found is None else found.group(1) - - -def body_of(request: Request) -> Mapping[str, JsonValue]: - try: - return _JSON_OBJECT.validate_json(request.body) - except ValueError: - return {} - - -def target_of(request: Request) -> str: - return unquote(request.target) - - -def answer(marker: str) -> str: - return f"answer marker-{marker}" - - -def reasoning_answer(marker: str) -> str: - return f"why marker-{marker} {answer(marker)}" - - -def _headers(marker: str) -> Mapping[str, str]: - return MappingProxyType({"x-amzn-requestid": marker}) - - -def _json_reply(status: int, payload: Mapping[str, JsonValue], marker: str) -> Reply: - return Reply(status=status, body=json.dumps(payload).encode(), headers=_headers(marker)) - - -def _error(status: int, message: str, marker: str) -> Reply: - return _json_reply(status, {"message": message}, marker) - - -def _effort_of(target: str, body: Mapping[str, JsonValue]) -> JsonValue: - if not _CONVERSE.match(target) and not _CONVERSE_STREAM.match(target): - return body.get("reasoning_effort") - fields: Final = body.get("additionalModelRequestFields") - reasoning: Final = fields.get("reasoning") if isinstance(fields, Mapping) else None - return reasoning.get("effort") if isinstance(reasoning, Mapping) else None - - -def forwarded_effort(request: Request) -> JsonValue: - return _effort_of(target_of(request), body_of(request)) - - -def _sse(frames: tuple[Mapping[str, JsonValue], ...], pause: float) -> Reply: - return Reply( - content_type="text/event-stream", - chunks=(*(b"data: " + json.dumps(frame).encode() + b"\n\n" for frame in frames), b"data: [DONE]\n\n"), - pause_between_chunks=pause, - ) - - -def _with_headers(reply: Reply, marker: str) -> Reply: - return Reply( - status=reply.status, - body=reply.body, - content_type=reply.content_type, - chunks=reply.chunks, - abort_after=reply.abort_after, - gate_after_first=reply.gate_after_first, - pause_between_chunks=reply.pause_between_chunks, - headers=_headers(marker), - ) - - -def _content_deltas(model: str, marker: str) -> tuple[str, ...]: - if "gpt-oss" in model: - return ("why ", f"marker-{marker}", " answer ", f"marker-{marker}") - return ("answer ", f"marker-{marker}") - - -def _chat_text(model: str, marker: str) -> str: - return reasoning_answer(marker) if "gpt-oss" in model else answer(marker) - - -def _chat_reply(model: str, marker: str, stream: bool, pause: float) -> Reply: - identity: Final = f"chatcmpl-{marker}" - if not stream: - return _json_reply( - 200, - { - "id": identity, - "object": "chat.completion", - "created": 1, - "model": model, - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": _chat_text(model, marker)}, - "finish_reason": "stop", - } - ], - "usage": dict(USAGE), - }, - marker, - ) - deltas: Final = _content_deltas(model, marker) - frames: Final = tuple( - { - "id": identity, - "object": "chat.completion.chunk", - "created": 1, - "model": model, - "choices": [{"index": 0, "delta": {"role": "assistant", "content": delta}, "finish_reason": None}], - } - for delta in deltas - ) - finish: Final[Mapping[str, JsonValue]] = { - "id": identity, - "object": "chat.completion.chunk", - "created": 1, - "model": model, - "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], - "usage": dict(USAGE), - } - return _with_headers(_sse((*frames, finish), pause), marker) - - -def _responses_reply(model: str, marker: str, stream: bool, pause: float) -> Reply: - identity: Final = f"resp_upstream_{marker}" - item_id: Final = f"msg_{marker}" - response: Final[Mapping[str, JsonValue]] = { - "id": identity, - "object": "response", - "created_at": 1, - "status": "completed", - "model": model, - "output": [ - { - "type": "message", - "id": item_id, - "status": "completed", - "role": "assistant", - "content": [{"type": "output_text", "text": answer(marker), "annotations": []}], - } - ], - "usage": {"input_tokens": 30, "output_tokens": 5, "total_tokens": 35}, - } - if not stream: - return _json_reply(200, response, marker) - events: Final[tuple[Mapping[str, JsonValue], ...]] = ( - { - "type": "response.created", - "sequence_number": 0, - "response": {**response, "status": "in_progress", "output": []}, - }, - { - "type": "response.output_text.delta", - "sequence_number": 1, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - "delta": answer(marker), - }, - {"type": "response.completed", "sequence_number": 2, "response": response}, - ) - return Reply( - content_type="text/event-stream", - chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events), - pause_between_chunks=pause, - headers=_headers(marker), - ) - - -def _converse_reply(marker: str) -> Reply: - return _json_reply( - 200, - { - "output": {"message": {"role": "assistant", "content": [{"text": answer(marker)}]}}, - "stopReason": "end_turn", - "usage": {"inputTokens": 9, "outputTokens": 5, "totalTokens": 14}, - "metrics": {"latencyMs": 1}, - }, - marker, - ) - - -def _converse_stream_reply(marker: str, pause: float) -> Reply: - events: Final[tuple[tuple[str, Mapping[str, JsonValue]], ...]] = ( - ("messageStart", {"role": "assistant"}), - ("contentBlockDelta", {"delta": {"text": "answer "}, "contentBlockIndex": 0}), - ("contentBlockDelta", {"delta": {"text": f"marker-{marker}"}, "contentBlockIndex": 0}), - ("contentBlockStop", {"contentBlockIndex": 0}), - ("messageStop", {"stopReason": "end_turn"}), - ("metadata", {"usage": {"inputTokens": 9, "outputTokens": 5, "totalTokens": 14}, "metrics": {"latencyMs": 1}}), - ) - return Reply( - content_type=EVENT_STREAM, - chunks=tuple(_aws_event_frame(kind, payload, "sc", marker) for kind, payload in events), - pause_between_chunks=pause, - headers=_headers(marker), - ) - - -def respond(request: Request, *, pause: float = 0.0) -> Reply: - target: Final = target_of(request) - marker: Final = marker_of(request) - if request.method == "GET": - if target == "/image.png": - return Reply(body=PNG_1X1, content_type="image/png", headers=_headers(marker)) - return _error(404, f"no scripted object at {target}", marker) - scripted_status: Final = _STATUS.search(request.body.decode(errors="replace")) - if scripted_status is not None: - status: Final = int(scripted_status.group(1)) - return _error(status, f"scripted {status}", marker) - body: Final = body_of(request) - effort: Final = _effort_of(target, body) - if effort is not None and (not isinstance(effort, str) or effort not in REASONING_EFFORTS): - return _error(400, f"Invalid reasoning effort: {json.dumps(effort)}", marker) - model: Final = str(body.get("model", "")) - stream: Final = body.get("stream") is True - if request.method == "POST" and target == NATIVE_CHAT: - return _chat_reply(model, marker, stream, pause) - if request.method == "POST" and target == NATIVE_RESPONSES: - return _responses_reply(model, marker, stream, pause) - if request.method == "POST" and _CONVERSE.match(target): - return _converse_reply(marker) - if request.method == "POST" and _CONVERSE_STREAM.match(target): - return _converse_stream_reply(marker, pause) - return _error(404, f"unknown bedrock route {request.method} {target}", marker) - - -def serve_peer(port: int, received: Synchronized[int], answer_first: int) -> None: - held: Final = threading.Event() - - def respond_or_hold(request: Request) -> Reply: - with received.get_lock(): - received.value += 1 - ordinal: Final = received.value - if ordinal > answer_first: - held.wait() - return respond(request) - - with wire_server(respond_or_hold, port=port): - threading.Event().wait() diff --git a/tests/integration/_support/openai_wire.py b/tests/integration/_support/openai_wire.py deleted file mode 100644 index bf9d5f194c1..00000000000 --- a/tests/integration/_support/openai_wire.py +++ /dev/null @@ -1,126 +0,0 @@ -from __future__ import annotations - -import json -from collections.abc import Callable -from typing import Final - -from integration._support.wire import Reply, Request, Wire -from pydantic import JsonValue - -from tests.integration._support.provider import MODEL_DISCOVERY - -_USAGE: Final = {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8} - - -def answering_model_discovery(respond: Callable[[Request], Reply]) -> Callable[[Request], Reply]: - def guarded(request: Request) -> Reply: - if (request.method, request.target) == MODEL_DISCOVERY: - return Reply(body=json.dumps({"object": "list", "data": []}).encode()) - return respond(request) - - return guarded - - -def posted_targets(wire: Wire) -> tuple[str, ...]: - return tuple(request.target for request in wire.drain() if request.method == "POST") - - -def openai_error(status: int) -> Reply: - return Reply( - status=status, - body=json.dumps({"error": {"message": f"scripted {status}", "type": "server_error", "code": None}}).encode(), - ) - - -def _data_frame(frame: dict[str, JsonValue]) -> bytes: - return b"data: " + json.dumps(frame).encode() + b"\n\n" - - -def _typed_frame(event: dict[str, JsonValue]) -> bytes: - return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() - - -def chat_stream(identity: str, model: str, text: str) -> tuple[bytes, bytes, bytes]: - chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": model} - role_only: Final = _data_frame({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}}]}) - content: Final = _data_frame({**chunk, "choices": [{"index": 0, "delta": {"content": text}}]}) - finish: Final = _data_frame( - {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": _USAGE} - ) - return (role_only, content, finish + b"data: [DONE]\n\n") - - -def chat_reply(identity: str, model: str, text: str, *, stream: bool) -> Reply: - if stream: - return Reply(content_type="text/event-stream", chunks=chat_stream(identity, model, text)) - return Reply( - body=json.dumps( - { - "id": identity, - "object": "chat.completion", - "created": 1, - "model": model, - "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], - "usage": _USAGE, - } - ).encode() - ) - - -def _response_object(identity: str, model: str, text: str) -> dict[str, JsonValue]: - return { - "id": identity, - "object": "response", - "created_at": 1, - "status": "completed", - "model": model, - "output": [_message_item(identity, text, "completed")], - "usage": {"input_tokens": 5, "output_tokens": 3, "total_tokens": 8}, - } - - -def _message_item(identity: str, text: str, status: str) -> dict[str, JsonValue]: - return { - "id": f"msg_{identity}", - "type": "message", - "role": "assistant", - "status": status, - "content": [{"type": "output_text", "text": text, "annotations": []}] if status == "completed" else [], - } - - -def responses_stream(identity: str, model: str, text: str) -> tuple[bytes, bytes, bytes]: - response: Final = _response_object(identity, model, text) - item: Final = _message_item(identity, text, "in_progress") - opened: Final = _typed_frame( - { - "type": "response.created", - "sequence_number": 0, - "response": {**response, "status": "in_progress", "output": []}, - } - ) + _typed_frame({"type": "response.output_item.added", "sequence_number": 1, "output_index": 0, "item": item}) - delta: Final = _typed_frame( - { - "type": "response.output_text.delta", - "sequence_number": 2, - "item_id": f"msg_{identity}", - "output_index": 0, - "content_index": 0, - "delta": text, - } - ) - closed: Final = _typed_frame( - { - "type": "response.output_item.done", - "sequence_number": 3, - "output_index": 0, - "item": _message_item(identity, text, "completed"), - } - ) + _typed_frame({"type": "response.completed", "sequence_number": 4, "response": response}) - return (opened, delta, closed) - - -def responses_reply(identity: str, model: str, text: str, *, stream: bool) -> Reply: - if stream: - return Reply(content_type="text/event-stream", chunks=responses_stream(identity, model, text)) - return Reply(body=json.dumps(_response_object(identity, model, text)).encode()) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 1b39eb81b01..1b137c38898 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -15,6 +15,7 @@ from redis import Redis from tests.integration._support.client import Gateway, eventually, gateway_from_environment from tests.integration._support.generation import LIFECYCLE_SETTINGS from tests.integration._support.manifest import OWNED_DIRECTORIES +from tests.integration._support.provider import SharedProvider, shared_provider from tests.integration._support.routing import RoutingPlugin COLLECTED: Final = pytest.StashKey[tuple[str, ...]]() @@ -126,6 +127,35 @@ def gateway() -> Iterator[Gateway]: yield value +@pytest.fixture(scope="session") +def shared_provider_server() -> Iterator[SharedProvider]: + if os.environ.get("PYTEST_XDIST_WORKER") and os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1": + pytest.fail("the shared fake provider needs tests to run one at a time; this group runs under pytest-xdist") + with shared_provider() as server: + yield server + late: Final = server.received() + assert late == (), ( + f"the shared fake provider got {[item.target for item in late]} after {server.last_test} finished" + ) + + +@pytest.fixture +def provider(shared_provider_server: SharedProvider, request: pytest.FixtureRequest) -> Iterator[SharedProvider]: + stray: Final = shared_provider_server.received() + shared_provider_server.replies.clear() + assert stray == (), ( + f"the shared fake provider got {[item.target for item in stray]} " + f"after {shared_provider_server.last_test} finished" + ) + yield shared_provider_server + shared_provider_server.last_test = request.node.nodeid + unused: Final = len(shared_provider_server.replies) + unread: Final = shared_provider_server.received() + shared_provider_server.replies.clear() + assert unused == 0, f"{unused} queued provider replies were never requested" + assert unread == (), f"the test never read the provider requests {[item.target for item in unread]}" + + @pytest.fixture def peer(gateway: Gateway) -> Iterator[Gateway]: url: Final = os.environ["INTEGRATION_PEER_URL"] diff --git a/tests/integration/providers/_count_tokens_system_lift.py b/tests/integration/providers/_count_tokens_system_lift.py deleted file mode 100644 index 56619d1c963..00000000000 --- a/tests/integration/providers/_count_tokens_system_lift.py +++ /dev/null @@ -1,159 +0,0 @@ -from collections.abc import Mapping -from dataclasses import dataclass -from types import MappingProxyType -from typing import Final - -import anthropic -import httpx -import openai -from integration._support.client import Gateway -from pydantic import JsonValue, TypeAdapter - -JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) -TOKEN_COUNTING_BETA: Final = "token-counting-2024-11-01" -ANTHROPIC_VERSION: Final = "2023-06-01" -REJECTION: Final = 'messages.0: Unexpected role "system". The Messages API accepts a top-level `system` parameter' - -USER_TEXT: Final = "Count this message" -USER: Final[dict[str, JsonValue]] = {"role": "user", "content": USER_TEXT} -ASSISTANT: Final[dict[str, JsonValue]] = {"role": "assistant", "content": "One."} -FOLLOW_UP: Final[dict[str, JsonValue]] = {"role": "user", "content": "Again"} -INSTRUCTION: Final = "You are a terse assistant" -REMINDER: Final = "Answer in one sentence" -LEADING: Final[dict[str, JsonValue]] = {"role": "system", "content": INSTRUCTION} -LIFTED: Final[dict[str, JsonValue]] = {"type": "text", "text": INSTRUCTION} -MID_SYSTEM: Final[dict[str, JsonValue]] = {"role": "system", "content": REMINDER} -EPHEMERAL: Final[dict[str, JsonValue]] = {"type": "ephemeral"} -ONE_HOUR: Final[dict[str, JsonValue]] = {"type": "ephemeral", "ttl": "1h"} -IMAGE_PART: Final[dict[str, JsonValue]] = { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" - }, -} -FIVE_KB: Final = "Answer in one sentence. " * 214 -CALLER_SYSTEM: Final = "Prefer metric units" -CALLER_BLOCKS: Final[list[JsonValue]] = [ - {"type": "text", "text": CALLER_SYSTEM}, - {"type": "text", "text": "Never guess", "cache_control": EPHEMERAL}, -] -TOOLS: Final[list[JsonValue]] = [ - { - "name": "get_weather", - "description": "Look up the current weather for a city", - "input_schema": { - "type": "object", - "properties": {"city": {"type": "string", "description": "City to look up"}}, - "required": ["city"], - }, - } -] - - -@dataclass(frozen=True, slots=True) -class LiftCase: - messages: tuple[dict[str, JsonValue], ...] - system: JsonValue | None - lifted: list[JsonValue] | None - - -LIFT_CASES: Final[Mapping[str, LiftCase]] = MappingProxyType( - { - "string": LiftCase((LEADING, USER), None, [LIFTED]), - "run_with_cache_control": LiftCase( - ( - {"role": "system", "content": INSTRUCTION, "cache_control": ONE_HOUR}, - { - "role": "system", - "content": [ - {"type": "text", "text": REMINDER, "cache_control": EPHEMERAL}, - {"type": "text", "text": ""}, - IMAGE_PART, - ], - }, - USER, - ), - None, - [ - {"type": "text", "text": INSTRUCTION, "cache_control": ONE_HOUR}, - {"type": "text", "text": REMINDER, "cache_control": EPHEMERAL}, - ], - ), - "caller_system_string_first": LiftCase( - (LEADING, USER), CALLER_SYSTEM, [{"type": "text", "text": CALLER_SYSTEM}, LIFTED] - ), - "caller_system_blocks_first": LiftCase((LEADING, USER), CALLER_BLOCKS, [*CALLER_BLOCKS, LIFTED]), - "caller_empty_system_dropped": LiftCase((LEADING, USER), "", [LIFTED]), - "empty_content_dropped": LiftCase(({"role": "system", "content": ""}, USER), None, None), - "image_only_content_dropped": LiftCase(({"role": "system", "content": [IMAGE_PART]}, USER), None, None), - "integer_content_dropped": LiftCase(({"role": "system", "content": 5}, USER), None, None), - "integer_text_part_dropped": LiftCase( - ({"role": "system", "content": [{"type": "text", "text": 7}]}, USER), None, None - ), - "five_kb": LiftCase(({"role": "system", "content": FIVE_KB}, USER), None, [{"type": "text", "text": FIVE_KB}]), - } -) -STRING_CASE: Final = LIFT_CASES["string"] - - -def count_request(model: str, case: LiftCase) -> dict[str, JsonValue]: - return {"model": model, "messages": list(case.messages), **({} if case.system is None else {"system": case.system})} - - -def expected_count_body(model: str, case: LiftCase) -> dict[str, JsonValue]: - conversation: Final[list[JsonValue]] = [message for message in case.messages if message["role"] != "system"] - return {"model": model, "messages": conversation, **({} if case.lifted is None else {"system": case.lifted})} - - -def _opening_message(message: JsonValue) -> bool: - return isinstance(message, dict) and message.get("role") in ("user", "assistant") - - -def _conversation_message(message: JsonValue) -> bool: - return isinstance(message, dict) and message.get("role") in ("user", "assistant", "system") - - -def _anthropic_tool(tool: JsonValue) -> bool: - return isinstance(tool, dict) and isinstance(tool.get("name"), str) and isinstance(tool.get("input_schema"), dict) - - -def accepts_count_body(body: Mapping[str, JsonValue]) -> bool: - # Anthropic, Azure AI Foundry and Bedrock Mantle count_tokens verdicts observed live on 2026-10-09: an empty - # messages list and a system role at messages[0] answer 400 invalid_request_error, a later system role counts - messages: Final = body.get("messages") - tools: Final = body.get("tools", []) - return ( - isinstance(messages, list) - and len(messages) > 0 - and _opening_message(messages[0]) - and all(map(_conversation_message, messages)) - and isinstance(body.get("system", ""), (str, list)) - and isinstance(tools, list) - and all(map(_anthropic_tool, tools)) - ) - - -def anthropic_client(gateway: Gateway) -> anthropic.Anthropic: - return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) - - -def async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic: - return anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) - - -def openai_client(gateway: Gateway) -> openai.OpenAI: - return openai.OpenAI( - base_url=f"{gateway.client.base_url}/v1", - api_key=gateway.key, - max_retries=0, - http_client=httpx.Client(trust_env=False), - ) - - -def async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI: - return openai.AsyncOpenAI( - base_url=f"{gateway.client.base_url}/v1", - api_key=gateway.key, - max_retries=0, - http_client=httpx.AsyncClient(trust_env=False), - ) 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 deleted file mode 100644 index b34b0b6b526..00000000000 --- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py +++ /dev/null @@ -1,581 +0,0 @@ -import asyncio -import json -import socket -import threading -import uuid -from collections.abc import Callable, Iterator, Mapping -from concurrent.futures import ThreadPoolExecutor -from contextlib import AbstractContextManager, ExitStack -from dataclasses import dataclass -from queue import SimpleQueue -from types import MappingProxyType -from typing import Final, cast - -import httpcore -import httpx -import psutil -import pytest -from anthropic.types import MessageParam -from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment -from integration._support.wire import Reply, Request, Wire, wire_server -from integration.providers._count_tokens_system_lift import ( - ANTHROPIC_VERSION, - ASSISTANT, - EPHEMERAL, - FOLLOW_UP, - IMAGE_PART, - INSTRUCTION, - JSON_OBJECT, - LEADING, - LIFT_CASES, - LIFTED, - MID_SYSTEM, - REJECTION, - REMINDER, - STRING_CASE, - TOKEN_COUNTING_BETA, - TOOLS, - USER, - USER_TEXT, - LiftCase, - accepts_count_body, - anthropic_client, - async_anthropic_client, - async_openai_client, - count_request, - expected_count_body, - openai_client, -) -from pydantic import JsonValue, TypeAdapter - -_MODEL: Final = "claude-opus-5-5" -_API_KEY: Final = "synthetic-count-tokens-key" -_COUNT: Final = 3131 -_PROVIDER_TOKENIZERS: Final = frozenset({"anthropic_api", "azure_ai_anthropic_api"}) -_SDK_MESSAGES: Final = cast( - list[MessageParam], [LEADING, USER] -) # cast-ok: the SDK types reject the role the proxy lifts -_STRING_BODY: Final = expected_count_body(_MODEL, STRING_CASE) -_PROXY_MODULE: Final = "integration._support.proxy" -_PROBES_PER_ROUND: Final = 8 -_CLIENT_ADDRESS: Final = TypeAdapter(tuple[str, int]) -_ARGUMENTS: Final = TypeAdapter(tuple[str, ...]) -_NAME: Final = TypeAdapter(str) - - -@dataclass(frozen=True, slots=True) -class _Provider: - prefix: str - target: str - tokenizer: str - azure: bool - - -@dataclass(frozen=True, slots=True) -class _Deployment: - provider: _Provider - port: int - model: str - - -_ANTHROPIC: Final = _Provider("anthropic", "/v1/messages/count_tokens", "anthropic_api", False) -_AZURE: Final = _Provider("azure_ai", "/anthropic/v1/messages/count_tokens", "azure_ai_anthropic_api", True) -_PROVIDERS: Final = MappingProxyType({"anthropic": _ANTHROPIC, "azure_ai": _AZURE}) - - -@pytest.fixture(params=_PROVIDERS.keys()) -def provider(request: pytest.FixtureRequest) -> _Provider: - name: Final[object] = request.param # pyright: ignore[reportAny] # pytest types the fixture param as Any - return _PROVIDERS[_NAME.validate_python(name)] - - -def _json_reply(status: int, payload: Mapping[str, JsonValue]) -> Reply: - return Reply(status=status, body=json.dumps(payload).encode()) - - -def _rejected(status: int) -> Reply: - return _json_reply(status, {"type": "error", "error": {"type": "invalid_request_error", "message": REJECTION}}) - - -def _counted(request: Request) -> Reply: - accepted: Final = accepts_count_body(JSON_OBJECT.validate_json(request.body)) - return _json_reply(200, {"input_tokens": _COUNT}) if accepted else _rejected(400) - - -def _rejecting(status: int) -> Callable[[Request], Reply]: - def count(_request: Request) -> Reply: - return _rejected(status) - - return count - - -def _holding(held: SimpleQueue[str], release: threading.Event, seconds: float) -> Callable[[Request], Reply]: - def hold(request: Request) -> Reply: - held.put(request.target) - assert release.wait(timeout=seconds), "Held count was never released" - return _counted(request) - - return hold - - -def _peer(provider: _Provider, count: Callable[[Request], Reply] = _counted) -> Callable[[Request], Reply]: - def respond(request: Request) -> Reply: - if request.target == provider.target: - return count(request) - return _json_reply(404, {"error": f"unscripted target {request.target}"}) - - return respond - - -def _listening(deployment: _Deployment, count: Callable[[Request], Reply] = _counted) -> AbstractContextManager[Wire]: - return wire_server(_peer(deployment.provider, count), port=deployment.port) - - -def _reserved_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return _CLIENT_ADDRESS.validate_python(reserve.getsockname())[1] - - -def _cmdline(process: psutil.Process) -> tuple[str, ...]: - try: - return _ARGUMENTS.validate_python(process.cmdline()) - except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): - return () - - -def _serves(cmdline: tuple[str, ...], proxy_port: int) -> bool: - if _PROXY_MODULE not in cmdline or "--port" not in cmdline: - return False - return cmdline[cmdline.index("--port") + 1] == str(proxy_port) - - -def _listens(process: psutil.Process, proxy_port: int) -> bool: - return any( - connection.status == psutil.CONN_LISTEN and connection.laddr and connection.laddr.port == proxy_port - for connection in process.net_connections(kind="tcp") - ) - - -def _proxy_workers(proxy_port: int) -> frozenset[int]: - (master,) = tuple(process for process in psutil.process_iter() if _serves(_cmdline(process), proxy_port)) - spawned: Final = frozenset(child.pid for child in master.children() if _listens(child, proxy_port)) - return spawned or frozenset({master.pid}) - - -def _holder(workers: frozenset[int], client_port: int) -> int | None: - def holds(pid: int) -> bool: - return any( - connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == client_port - for connection in psutil.Process(pid).net_connections(kind="tcp") - ) - - return next((pid for pid in sorted(workers) if holds(pid)), None) - - -def _client_port(response: httpx.Response) -> int: - stream: Final[object] = response.extensions["network_stream"] # pyright: ignore[reportAny] # httpx types extensions as Any - assert isinstance(stream, httpcore.NetworkStream), stream - return _CLIENT_ADDRESS.validate_python(stream.get_extra_info("client_addr"))[1] - - -def _probe(gateway: Gateway, body: Mapping[str, JsonValue], workers: frozenset[int]) -> tuple[int | None, JsonValue]: - with httpx.Client(base_url=str(gateway.client.base_url), timeout=30, trust_env=False) as client: - response: Final = client.post( - "/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {gateway.key}"} - ) - return _holder(workers, _client_port(response)), JSON_OBJECT.validate_json(response.content).get("input_tokens") - - -def _round( - gateway: Gateway, body: Mapping[str, JsonValue], workers: frozenset[int] -) -> tuple[tuple[int | None, JsonValue], ...]: - def probe(_index: int) -> tuple[int | None, JsonValue]: - return _probe(gateway, body, workers) - - with ThreadPoolExecutor(max_workers=_PROBES_PER_ROUND) as pool: - return tuple(pool.map(probe, range(_PROBES_PER_ROUND))) - - -def _settled_on_every_worker(gateway: Gateway, body: Mapping[str, JsonValue]) -> None: - proxy_port: Final = gateway.client.base_url.port - assert proxy_port is not None, gateway.client.base_url - workers: Final = _proxy_workers(proxy_port) - eventually( - lambda: _round(gateway, body, workers), - lambda observed: ( - frozenset(pid for pid, _ in observed) == workers and all(count == _COUNT for _, count in observed) - ), - seconds=60, - ) - - -def _deploy(gateway: Gateway, scenario: Scenario, provider: _Provider) -> _Deployment: - port: Final = _reserved_port() - model: Final = scenario.model( - model_info=None, model=f"{provider.prefix}/{_MODEL}", api_base=f"http://127.0.0.1:{port}", api_key=_API_KEY - ) - with wire_server(_peer(provider), port=port): - _settled_on_every_worker(gateway, {"model": model, "messages": [USER]}) - return _Deployment(provider, port, model) - - -@pytest.fixture(scope="module") -def deployments() -> Iterator[Mapping[str, _Deployment]]: - with gateway_from_environment() as gateway, gateway.scenario() as scenario: - yield MappingProxyType({chosen.prefix: _deploy(gateway, scenario, chosen) for chosen in (_ANTHROPIC, _AZURE)}) - - -@pytest.fixture -def deployment(provider: _Provider, deployments: Mapping[str, _Deployment]) -> _Deployment: - return deployments[provider.prefix] - - -def _count(gateway: Gateway, body: Mapping[str, JsonValue], key: str | None = None) -> httpx.Response: - return gateway.request("POST", "/v1/messages/count_tokens", body, key=key) - - -def _payload(response: httpx.Response) -> dict[str, JsonValue]: - assert response.status_code == 200, response.text - return JSON_OBJECT.validate_json(response.content) - - -def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int: - response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "false"}) - payload: Final = _payload(response) - total: Final = payload["total_tokens"] - assert payload["tokenizer_type"] not in _PROVIDER_TOKENIZERS, response.text - assert isinstance(total, int) and total > 0 and total != _COUNT, response.text - return total - - -def _bodies(wire: Wire, provider: _Provider) -> tuple[dict[str, JsonValue], ...]: - received: Final = wire.drain() - for request in received: - assert (request.method, request.target) == ("POST", provider.target), request.target - assert request.headers["anthropic-version"] == ANTHROPIC_VERSION, request.headers - assert TOKEN_COUNTING_BETA in request.headers["anthropic-beta"], request.headers - assert request.headers["content-type"] == "application/json", request.headers - assert request.headers["x-api-key"] == _API_KEY, request.headers - assert (request.headers.get("api-key") == _API_KEY) is provider.azure, request.headers - return tuple(JSON_OBJECT.validate_json(request.body) for request in received) - - -def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]: - return tuple( - stack.enter_context(httpx.Client(base_url=base_url, timeout=30, trust_env=False)) for _ in range(count) - ) - - -def _counted_on(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue]: - response: Final = client.post( - "/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {key}"} - ) - return response.status_code, JSON_OBJECT.validate_json(response.content).get("input_tokens") - - -@pytest.mark.parametrize("case", LIFT_CASES.values(), ids=LIFT_CASES.keys()) -def test_messages_count_tokens_lifts_the_leading_system_run( - deployment: _Deployment, gateway: Gateway, case: LiftCase -) -> None: - with _listening(deployment) as peer: - response: Final = _count(gateway, count_request(deployment.model, case)) - assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text - assert _bodies(peer, deployment.provider) == (expected_count_body(_MODEL, case),) - - -def test_anthropic_sdk_count_tokens_lifts_the_leading_system_message(deployment: _Deployment, gateway: Gateway) -> None: - with _listening(deployment) as peer: - counted: Final = anthropic_client(gateway).messages.count_tokens(model=deployment.model, messages=_SDK_MESSAGES) - assert counted.input_tokens == _COUNT, counted - assert _bodies(peer, deployment.provider) == (_STRING_BODY,) - - -def test_async_anthropic_sdk_count_tokens_lifts_the_leading_system_message( - deployment: _Deployment, gateway: Gateway -) -> None: - with _listening(deployment) as peer: - counted: Final = asyncio.run( - async_anthropic_client(gateway).messages.count_tokens(model=deployment.model, messages=_SDK_MESSAGES) - ) - assert counted.input_tokens == _COUNT, counted - assert _bodies(peer, deployment.provider) == (_STRING_BODY,) - - -def test_utils_token_counter_call_endpoint_counts_a_leading_system_through_the_provider( - deployment: _Deployment, gateway: Gateway -) -> None: - with _listening(deployment) as peer: - response: Final = gateway.request( - "POST", - "/utils/token_counter", - count_request(deployment.model, STRING_CASE), - params={"call_endpoint": "true"}, - ) - payload: Final = _payload(response) - expected_tokenizer: Final = deployment.provider.tokenizer - assert (payload["total_tokens"], payload["tokenizer_type"]) == (_COUNT, expected_tokenizer), response.text - assert payload["original_response"] == {"input_tokens": _COUNT}, response.text - assert (payload["request_model"], payload["model_used"]) == (deployment.model, _MODEL), response.text - assert _bodies(peer, deployment.provider) == (_STRING_BODY,) - - -def test_utils_token_counter_local_mode_never_calls_the_peer_for_a_leading_system( - deployment: _Deployment, gateway: Gateway -) -> None: - with _listening(deployment) as peer: - assert _local_count(gateway, count_request(deployment.model, STRING_CASE)) > 0 - assert peer.drain() == () - - -def test_responses_input_tokens_lifts_instructions(deployment: _Deployment, gateway: Gateway) -> None: - with _listening(deployment) as peer: - response: Final = gateway.request( - "POST", - "/v1/responses/input_tokens", - {"model": deployment.model, "input": USER_TEXT, "instructions": INSTRUCTION}, - ) - assert _payload(response) == { - "object": "response.input_tokens", - "input_tokens": _COUNT, - "litellm_count_source": "provider", - }, response.text - assert _bodies(peer, deployment.provider) == (_STRING_BODY,) - - -def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_item( - deployment: _Deployment, gateway: Gateway -) -> None: - with _listening(deployment) as peer: - response: Final = gateway.request( - "POST", - "/v1/responses/input_tokens", - {"model": deployment.model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION}, - ) - assert _payload(response) == { - "object": "response.input_tokens", - "input_tokens": _COUNT, - "litellm_count_source": "provider", - }, response.text - assert _bodies(peer, deployment.provider) == ( - {"model": _MODEL, "messages": [USER], "system": [LIFTED, {"type": "text", "text": REMINDER}]}, - ) - - -def test_openai_sdk_input_tokens_lifts_instructions(deployment: _Deployment, gateway: Gateway) -> None: - with _listening(deployment) as peer: - counted: Final = openai_client(gateway).responses.input_tokens.count( - model=deployment.model, input=USER_TEXT, instructions=INSTRUCTION - ) - assert counted.input_tokens == _COUNT, counted - assert _bodies(peer, deployment.provider) == (_STRING_BODY,) - - -def test_async_openai_sdk_input_tokens_lifts_instructions(deployment: _Deployment, gateway: Gateway) -> None: - with _listening(deployment) as peer: - counted: Final = asyncio.run( - async_openai_client(gateway).responses.input_tokens.count( - model=deployment.model, input=USER_TEXT, instructions=INSTRUCTION - ) - ) - assert counted.input_tokens == _COUNT, counted - assert _bodies(peer, deployment.provider) == (_STRING_BODY,) - - -def test_messages_count_tokens_forwards_tools_beside_the_lifted_system( - deployment: _Deployment, gateway: Gateway -) -> None: - with _listening(deployment) as peer: - response: Final = _count(gateway, {**count_request(deployment.model, STRING_CASE), "tools": TOOLS}) - assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text - assert _bodies(peer, deployment.provider) == ({**_STRING_BODY, "tools": TOOLS},) - - -def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place( - deployment: _Deployment, gateway: Gateway -) -> None: - with _listening(deployment) as peer: - body: Final[dict[str, JsonValue]] = { - "model": deployment.model, - "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], - } - response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text - assert _bodies(peer, deployment.provider) == ( - {"model": _MODEL, "messages": [USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], "system": [LIFTED]}, - ) - - -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]} - response: Final = _count(gateway, body) - 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_surfaces_the_peer_error_for_a_non_text_system( - deployment: _Deployment, gateway: Gateway -) -> None: - with _listening(deployment) as peer: - response: Final = _count(gateway, {**count_request(deployment.model, STRING_CASE), "system": 5}) - 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},) - - -def test_messages_count_tokens_repeated_request_lifts_each_time(deployment: _Deployment, gateway: Gateway) -> None: - with _listening(deployment) as peer: - answers: Final = tuple( - _payload(_count(gateway, count_request(deployment.model, STRING_CASE))) for _ in range(2) - ) - assert answers == ({"input_tokens": _COUNT, "litellm_count_source": "provider"},) * 2 - assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * 2 - - -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]} - response: Final = _count(gateway, body) - 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, "litellm_count_source": "provider"}, follow_up.text - assert _bodies(peer, deployment.provider) == (_STRING_BODY,) - - -def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value( - deployment: _Deployment, gateway: Gateway -) -> None: - with _listening(deployment) as peer: - first: Final = json.dumps([USER]) - last: Final = json.dumps([LEADING, USER]) - response: Final = gateway.client.post( - "/v1/messages/count_tokens", - content=f'{{"model": "{deployment.model}", "messages": {first}, "messages": {last}}}', - headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, - ) - assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text - assert _bodies(peer, deployment.provider) == (_STRING_BODY,) - - -@pytest.mark.parametrize("status", [400, 403, 404, 500, 503]) -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) - response: Final = _count(gateway, body) - assert response.status_code == status, response.text - assert "input_tokens" not in response.text, response.text - assert _bodies(peer, deployment.provider) == (_STRING_BODY,) - - -def test_messages_count_tokens_unauthenticated_request_never_reaches_the_peer( - deployment: _Deployment, gateway: Gateway -) -> None: - with _listening(deployment) as peer: - response: Final = _count(gateway, count_request(deployment.model, STRING_CASE), key="sk-not-a-key") - assert response.status_code == 401, response.text - assert peer.drain() == () - - -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) - - def count(client: httpx.Client) -> tuple[int, JsonValue]: - return _counted_on(client, gateway.key, body) - - with _listening(deployment) as peer: - assert tuple(pool.map(count, clients)) == ((200, _COUNT),) * len(clients) - assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * 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) - - -def test_slow_peer_holds_concurrent_lifted_counts_without_stalling_the_proxy( - deployment: _Deployment, gateway: Gateway -) -> None: - held: Final[SimpleQueue[str]] = SimpleQueue() - release: Final = threading.Event() - with ExitStack() as stack: - clients: Final = _clients(stack, str(gateway.client.base_url), 6) - peer: Final = stack.enter_context(_listening(deployment, _holding(held, release, 20))) - pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) - stack.callback(release.set) - body: Final = count_request(deployment.model, STRING_CASE) - futures: Final = tuple(pool.submit(_counted_on, client, gateway.key, body) for client in clients) - eventually(held.qsize, lambda size: size == len(clients), seconds=30) - assert gateway.request("GET", "/health/liveliness").status_code == 200 - assert _local_count(gateway, body) > 0 - assert not any(future.done() for future in futures) - release.set() - assert tuple(future.result(timeout=30) for future in futures) == ((200, _COUNT),) * len(clients) - assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * len(clients) - - -def test_chat_completions_on_the_same_deployment_keeps_a_mid_conversation_system_in_place( - deployments: Mapping[str, _Deployment], gateway: Gateway -) -> None: - identity: Final = f"msg_{uuid.uuid4().hex}" - anthropic_deployment: Final = deployments[_ANTHROPIC.prefix] - - def respond(request: Request) -> Reply: - assert (request.method, request.target) == ("POST", "/v1/messages"), request.target - return _json_reply( - 200, - { - "id": identity, - "type": "message", - "role": "assistant", - "model": _MODEL, - "content": [{"type": "text", "text": "done"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 12, "output_tokens": 3}, - }, - ) - - with wire_server(respond, port=anthropic_deployment.port) as peer: - reminder: Final[dict[str, JsonValue]] = { - "role": "system", - "content": [ - {"type": "text", "text": REMINDER, "cache_control": EPHEMERAL}, - {"type": "text", "text": ""}, - IMAGE_PART, - ], - } - response: Final = gateway.request( - "POST", - "/v1/chat/completions", - { - "model": anthropic_deployment.model, - "max_tokens": 16, - "messages": [LEADING, USER, reminder, ASSISTANT, FOLLOW_UP], - }, - ) - assert response.status_code == 200, response.text - (sent,) = peer.drain() - body: Final = JSON_OBJECT.validate_json(sent.body) - assert body["system"] == [LIFTED], body - assert body["messages"] == [ - {"role": "user", "content": [{"type": "text", "text": USER_TEXT}]}, - {"role": "system", "content": [{"type": "text", "text": REMINDER, "cache_control": EPHEMERAL}]}, - {"role": "assistant", "content": [{"type": "text", "text": "One."}]}, - {"role": "user", "content": [{"type": "text", "text": "Again"}]}, - ], body diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py deleted file mode 100644 index e20a397a708..00000000000 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ /dev/null @@ -1,1568 +0,0 @@ -import asyncio -import json -import re -import signal -import socket -import threading -import uuid -from collections.abc import Callable, Iterator, Mapping, Sequence -from concurrent.futures import Future, ThreadPoolExecutor -from contextlib import ExitStack -from pathlib import Path -from queue import SimpleQueue -from types import MappingProxyType -from typing import Final, cast - -import anthropic -import httpx -import psutil -import pytest -import yaml -from anthropic.types import MessageCountTokensToolParam, MessageParam -from integration._support.bedrock_runtime_peer import answer, marker_of, target_of -from integration._support.bedrock_runtime_peer import respond as runtime_generation -from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value -from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process -from integration._support.wire import Reply, Request, Wire, wire_server -from integration.providers._count_tokens_system_lift import ( - ASSISTANT, - FOLLOW_UP, - INSTRUCTION, - LEADING, - LIFT_CASES, - LIFTED, - MID_SYSTEM, - REMINDER, - STRING_CASE, - USER, - USER_TEXT, - LiftCase, - accepts_count_body, - async_openai_client, - count_request, - expected_count_body, - openai_client, -) -from pydantic import JsonValue, TypeAdapter - -pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) - -_OPUS: Final = "global.anthropic.claude-opus-4-8" -_OPUS_BASE: Final = "anthropic.claude-opus-4-8" -_SONNET: Final = "anthropic.claude-sonnet-4-6" -_NOVA: Final = "amazon.nova-lite-v1:0" -_REGION: Final = "us-east-1" -_OWNED_OPUS: Final = "mantle-opus" -_OWNED_SONNET: Final = "mantle-sonnet" -_OWNED_NOVA: Final = "mantle-nova" -_OWNED_MODELS: Final = MappingProxyType({_OWNED_OPUS: _OPUS, _OWNED_SONNET: _SONNET, _OWNED_NOVA: _NOVA}) -_ACCESS_KEY: Final = "AKIAINTEGRATIONMANTLE" -_SECRET_KEY: Final = "integration-mantle-secret" -_SIGV4_SCOPE: Final = f"/{_REGION}/bedrock/aws4_request" -_MANTLE_TARGET: Final = "/anthropic/v1/messages/count_tokens" -_MANTLE_VERSION: Final = "2023-06-01" -_MANTLE_COUNT: Final = 4242 -_RUNTIME_COUNT: Final = 1345 -_UNSUPPORTED: Final = "The provided model doesn't support counting tokens." -_REJECTION: Final = "scripted mantle rejection" -_COUNT_TARGET: Final = re.compile(r"^/model/(.+)/count-tokens$") -_INVOKE_TARGET: Final = re.compile(r"^/model/(.+)/invoke$") -_SCRIPTED_STATUS: Final = re.compile(r"status=(\d{3})") -_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") -_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) -_JSON_LIST: Final = TypeAdapter(list[JsonValue]) - -_SDK_MESSAGES: Final[list[MessageParam]] = [{"role": "user", "content": "Count this message"}] -_MESSAGES: Final = _JSON_LIST.validate_json(json.dumps(_SDK_MESSAGES)) -_SYSTEM: Final = "You are a terse assistant that answers in one sentence" -_SYSTEM_BLOCKS: Final[list[JsonValue]] = [ - {"type": "text", "text": "You are a terse assistant"}, - {"type": "text", "text": "Answer in one sentence"}, -] -_SDK_TOOLS: Final[list[MessageCountTokensToolParam]] = [ - { - "name": "get_weather", - "description": "Look up the current weather for a city", - "input_schema": { - "type": "object", - "properties": {"city": {"type": "string", "description": "City to look up"}}, - "required": ["city"], - }, - } -] -_TOOLS: Final = _JSON_LIST.validate_json(json.dumps(_SDK_TOOLS)) -_GEMINI_BODY: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "Count this"}]}]} -_GEMINI_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this"}] - - -def _mantle_body(**fields: JsonValue) -> dict[str, JsonValue]: - return {"model": _OPUS_BASE, "messages": _MESSAGES, **fields} - - -_MANTLE_BARE: Final = _mantle_body() -_MANTLE_LIFTED: Final = _mantle_body(system=[LIFTED]) -_LIFT_SDK_MESSAGES: Final = cast( - list[MessageParam], [LEADING, USER] -) # cast-ok: the SDK types reject the role the proxy lifts -_MANTLE_FULL: Final = _mantle_body(system=_SYSTEM, tools=_TOOLS) - - -def _json_reply(status: int, payload: Mapping[str, JsonValue]) -> Reply: - return Reply(status=status, body=json.dumps(payload).encode()) - - -def _runtime_count(request: Request, model: str) -> Reply: - scripted: Final = _SCRIPTED_STATUS.search(request.body.decode(errors="replace")) - if scripted is not None: - return _json_reply(int(scripted.group(1)), {"message": f"scripted {scripted.group(1)}"}) - if "sonnet" in model: - return _json_reply(200, {"inputTokens": _RUNTIME_COUNT}) - return _json_reply(400, {"message": _UNSUPPORTED}) - - -def _invoke_reply(marker: str) -> Reply: - return _json_reply( - 200, - { - "id": f"msg_{marker}", - "type": "message", - "role": "assistant", - "model": _OPUS_BASE, - "content": [{"type": "text", "text": answer(marker)}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 5, "output_tokens": 3}, - }, - ) - - -def _runtime(request: Request) -> Reply: - target: Final = target_of(request) - counted: Final = _COUNT_TARGET.match(target) - if counted is not None: - return _runtime_count(request, counted.group(1)) - if _INVOKE_TARGET.match(target): - return _invoke_reply(marker_of(request)) - return runtime_generation(request) - - -def _mantle_counted(_request: Request) -> Reply: - return _json_reply(200, {"input_tokens": _MANTLE_COUNT}) - - -def _rejected(status: int) -> Reply: - return _json_reply(status, {"type": "error", "error": {"type": "invalid_request_error", "message": _REJECTION}}) - - -def _rejecting(status: int) -> Callable[[Request], Reply]: - def count(_request: Request) -> Reply: - return _rejected(status) - - return count - - -def _strict(request: Request) -> Reply: - accepted: Final = accepts_count_body(_JSON_OBJECT.validate_json(request.body)) - return _mantle_counted(request) if accepted else _rejected(400) - - -def _mantle(count: Callable[[Request], Reply] = _mantle_counted) -> Callable[[Request], Reply]: - def respond(request: Request) -> Reply: - if request.target == _MANTLE_TARGET: - return count(request) - return _json_reply(404, {"error": f"unscripted mantle target {request.target}"}) - - return respond - - -def _mantle_environment(port: int) -> Mapping[str, str]: - return {"BEDROCK_MANTLE_API_BASE": f"http://127.0.0.1:{port}"} - - -_INHERITED_BEARER: Final = ("AWS_BEARER_TOKEN_BEDROCK",) - - -def _reserved_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return reserve.getsockname()[1] - - -def _closed_port_url() -> str: - return f"http://127.0.0.1:{_reserved_port()}" - - -@pytest.fixture(scope="module") -def mantle_port() -> int: - return _reserved_port() - - -@pytest.fixture(scope="module") -def counting_proxy(tmp_path_factory: pytest.TempPathFactory, mantle_port: int) -> Iterator[OwnedProxy]: - with ( - gateway_from_environment() as gateway, - owned_proxy_process( - gateway, - tmp_path_factory.mktemp("mantle-count"), - _mantle_environment(mantle_port), - workers=1, - remove_environment=_INHERITED_BEARER, - ) as owned, - ): - yield owned - - -def _litellm_params(model: str, api_base: str) -> dict[str, JsonValue]: - return { - "model": f"bedrock/{model}", - "api_base": api_base, - "aws_access_key_id": _ACCESS_KEY, - "aws_secret_access_key": _SECRET_KEY, - "aws_region_name": _REGION, - } - - -def _owned_config(path: Path, runtime_url: str, settings: Mapping[str, JsonValue]) -> Path: - config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) - litellm_settings: Final = object_value(config["litellm_settings"]) - path.write_text( - yaml.safe_dump( - { - **config, - "model_list": [ - {"model_name": name, "litellm_params": _litellm_params(model, runtime_url)} - for name, model in _OWNED_MODELS.items() - ], - "litellm_settings": {**litellm_settings, **settings}, - } - ) - ) - return path - - -def _deployment(scenario: Scenario, api_base: str, model: str = _OPUS) -> str: - return scenario.model(model_info=None, api_key=None, **_litellm_params(model, api_base)) - - -def _bare(model: str) -> dict[str, JsonValue]: - return {"model": model, "messages": _MESSAGES} - - -def _full(model: str) -> dict[str, JsonValue]: - return {**_bare(model), "system": _SYSTEM, "tools": _TOOLS} - - -def _count(gateway: Gateway, body: Mapping[str, JsonValue]) -> httpx.Response: - return gateway.request("POST", "/v1/messages/count_tokens", body) - - -def _payload(response: httpx.Response) -> dict[str, JsonValue]: - assert response.status_code == 200, response.text - return _JSON_OBJECT.validate_json(response.content) - - -def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int: - response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "false"}) - payload: Final = _payload(response) - total: Final = payload["total_tokens"] - assert payload["tokenizer_type"] not in ("bedrock_api", "bedrock_mantle_api"), response.text - assert isinstance(total, int) and total > 0 and total not in (_MANTLE_COUNT, _RUNTIME_COUNT), response.text - return total - - -def _assert_sigv4(request: Request) -> None: - authorization: Final = request.headers.get("authorization", "") - assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={_ACCESS_KEY}/"), request.headers - assert _SIGV4_SCOPE in authorization, authorization - - -def _mantle_requests(requests: Sequence[Request]) -> tuple[dict[str, JsonValue], ...]: - for request in requests: - assert (request.method, request.target) == ("POST", _MANTLE_TARGET), request.target - assert request.headers["anthropic-version"] == _MANTLE_VERSION, request.headers - assert request.headers["content-type"] == "application/json", request.headers - _assert_sigv4(request) - return tuple(_JSON_OBJECT.validate_json(request.body) for request in requests) - - -def _mantle_bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]: - return _mantle_requests(wire.drain()) - - -def _runtime_count_targets(wire: Wire) -> tuple[str, ...]: - counts: Final = tuple(request for request in wire.drain() if _COUNT_TARGET.match(target_of(request))) - for request in counts: - assert request.method == "POST", request.method - _assert_sigv4(request) - return tuple(target_of(request) for request in counts) - - -def _clients(stack: ExitStack, base_url: str, count: int, timeout: float = 30) -> tuple[httpx.Client, ...]: - return tuple( - stack.enter_context(httpx.Client(base_url=base_url, timeout=timeout, trust_env=False)) for _ in range(count) - ) - - -def _counted_on(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue]: - response: Final = client.post( - "/v1/messages/count_tokens", json=dict(body), headers={"Authorization": f"Bearer {key}"} - ) - return response.status_code, _JSON_OBJECT.validate_json(response.content).get("input_tokens") - - -def _generated_then_counted( - client: httpx.Client, key: str, model: str, body: Mapping[str, JsonValue] -) -> tuple[int, int, JsonValue]: - generated: Final = client.post( - "/v1/messages", - json={ - "model": model, - "max_tokens": 16, - "messages": [{"role": "user", "content": f"Generate before counting marker-{uuid.uuid4().hex}"}], - }, - headers={"Authorization": f"Bearer {key}"}, - ) - return generated.status_code, *_counted_on(client, key, body) - - -def _counted_or_dropped(client: httpx.Client, key: str, body: Mapping[str, JsonValue]) -> tuple[int, JsonValue] | None: - try: - return _counted_on(client, key, body) - except httpx.TransportError: - return None - - -def _probed_then_counted_or_dropped( - client: httpx.Client, key: str, body: Mapping[str, JsonValue], probed: SimpleQueue[int] -) -> tuple[int, tuple[int, JsonValue] | None]: - port: Final = _local_port(client) - probed.put(port) - return port, _counted_or_dropped(client, key, body) - - -def _local_port(client: httpx.Client) -> int: - with client.stream("GET", "/health/liveliness") as response: - port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1]) - response.read() - assert response.status_code == 200, response.text - return port - - -def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]: - return frozenset( - connection.raddr.port - for connection in psutil.Process(pid).net_connections(kind="tcp") - if connection.raddr and connection.laddr.port == proxy_port - ) - - -def _holding(held: SimpleQueue[str], release: threading.Event, seconds: float) -> Callable[[Request], Reply]: - def hold(request: Request) -> Reply: - held.put(request.target) - assert release.wait(timeout=seconds), "Held count was never released" - return _mantle_counted(request) - - return hold - - -def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic: - return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) - - -def _async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic: - return anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) - - -@pytest.mark.parametrize("system", [_SYSTEM, _SYSTEM_BLOCKS], ids=["string", "blocks"]) -def test_messages_count_tokens_counts_through_mantle_when_the_runtime_cannot( - counting_proxy: OwnedProxy, mantle_port: int, system: JsonValue -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text - assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) - assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) - - -def test_messages_count_tokens_without_system_or_tools_sends_a_bare_body_to_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_MANTLE_BARE,) - - -def test_utils_token_counter_call_endpoint_reports_the_mantle_tokenizer( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request( - "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} - ) - payload: Final = _payload(response) - assert (payload["total_tokens"], payload["tokenizer_type"]) == (_MANTLE_COUNT, "bedrock_mantle_api") - assert payload["original_response"] == {"input_tokens": _MANTLE_COUNT}, response.text - assert (payload["request_model"], payload["model_used"]) == (model, _OPUS), response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_MANTLE_FULL,) - - -def test_gemini_count_tokens_route_reports_the_mantle_total(counting_proxy: OwnedProxy, mantle_port: int) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) - assert _payload(response) == { - "totalTokens": _MANTLE_COUNT, - "promptTokensDetails": [], - "litellm_count_source": "provider", - }, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == ({"model": _OPUS_BASE, "messages": _GEMINI_MESSAGES},) - - -def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock_counts( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url, _SONNET) - response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) - assert _payload(response) == { - "totalTokens": _RUNTIME_COUNT, - "promptTokensDetails": [], - "litellm_count_source": "provider", - }, response.text - assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) - assert mantle.drain() == () - - -def test_runtime_404_on_claude_falls_back_to_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: - gateway: Final = counting_proxy.gateway - system: Final = f"{_SYSTEM} status=404" - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) - - -def test_runtime_404_and_mantle_404_surface_provider_error(counting_proxy: OwnedProxy, mantle_port: int) -> None: - gateway: Final = counting_proxy.gateway - system: Final = f"{_SYSTEM} status=404" - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_rejecting(404)), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) - assert response.status_code == 404, response.text - assert _REJECTION in response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) - - -def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url, _NOVA) - body: Final = {**_bare(model), "messages": [{"role": "user", "content": "status=404"}]} - response: Final = _count(gateway, body) - assert response.status_code == 404, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert mantle.drain() == () - - -def test_google_count_tokens_forwards_system_and_tools_to_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - system_text: Final = "Use the provided tool" - tool_schema: Final = {"type": "object", "properties": {"location": {"type": "string"}}} - body: Final[dict[str, JsonValue]] = { - "contents": [{"role": "user", "parts": [{"text": "status=404"}]}], - "systemInstruction": {"parts": [{"text": system_text}]}, - "tools": [ - { - "functionDeclarations": [ - {"name": "get_weather", "description": "Look up weather", "parameters": tool_schema} - ] - } - ], - } - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) - assert _payload(response) == { - "totalTokens": _MANTLE_COUNT, - "promptTokensDetails": [], - "litellm_count_source": "provider", - }, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == ( - _mantle_body( - messages=[{"role": "user", "content": "status=404"}], - system=[{"type": "text", "text": system_text}], - tools=[ - { - "name": "get_weather", - "description": "Look up weather", - "input_schema": tool_schema, - } - ], - ), - ) - - -def test_google_count_tokens_normalizes_uppercase_parameters_for_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - body: Final[dict[str, JsonValue]] = { - "contents": [{"role": "user", "parts": [{"text": "status=404"}]}], - "tools": [ - { - "functionDeclarations": [ - { - "name": "get_weather", - "parameters": { - "type": "OBJECT", - "properties": {"location": {"type": "STRING"}}, - "required": ["location"], - }, - } - ] - } - ], - } - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) - assert _payload(response) == { - "totalTokens": _MANTLE_COUNT, - "promptTokensDetails": [], - "litellm_count_source": "provider", - }, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == ( - _mantle_body( - messages=[{"role": "user", "content": "status=404"}], - tools=[ - { - "name": "get_weather", - "input_schema": { - "type": "object", - "properties": {"location": {"type": "string"}}, - "required": ["location"], - }, - } - ], - ), - ) - - -def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request( - "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"} - ) - assert _payload(response) == { - "object": "response.input_tokens", - "input_tokens": _MANTLE_COUNT, - "litellm_count_source": "provider", - }, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_MANTLE_BARE,) - - -def test_anthropic_sdk_count_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - counted: Final = _anthropic_client(gateway).messages.count_tokens( - model=model, messages=_SDK_MESSAGES, system=_SYSTEM, tools=_SDK_TOOLS - ) - assert counted.input_tokens == _MANTLE_COUNT, counted - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_MANTLE_FULL,) - - -def test_async_anthropic_sdk_count_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - counted: Final = asyncio.run( - _async_anthropic_client(gateway).messages.count_tokens(model=model, messages=_SDK_MESSAGES, system=_SYSTEM) - ) - assert counted.input_tokens == _MANTLE_COUNT, counted - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_mantle_body(system=_SYSTEM),) - - -def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url, _SONNET) - assert _payload(_count(gateway, _full(model))) == { - "input_tokens": _RUNTIME_COUNT, - "litellm_count_source": "provider", - } - detailed: Final = gateway.request( - "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} - ) - payload: Final = _payload(detailed) - assert (payload["total_tokens"], payload["tokenizer_type"]) == (_RUNTIME_COUNT, "bedrock_api"), detailed.text - assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) * 2 - assert mantle.drain() == () - - -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 - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url, _NOVA) - response: Final = _count(gateway, _bare(model)) - 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() == () - - -def test_runtime_403_on_a_claude_model_never_reaches_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - body: Final[dict[str, JsonValue]] = { - "model": model, - "messages": [{"role": "user", "content": "status=403 Count this message"}], - } - response: Final = _count(gateway, body) - 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_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 500 <= response.status_code < 600, response.text - assert "input_tokens" not in response.text, response.text - assert mantle.drain() == () - - -def test_bedrock_passthrough_count_tokens_still_answers_the_runtime_rejection( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request("POST", "/bedrock/v1/messages/count_tokens", _full(model)) - assert response.status_code == 400, response.text - assert _UNSUPPORTED in response.text and "input_tokens" not in response.text, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert mantle.drain() == () - - -def test_responses_input_tokens_with_instructions_counts_through_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request( - "POST", - "/v1/responses/input_tokens", - {"model": model, "input": "Count this message", "instructions": "Be terse"}, - ) - assert _payload(response) == { - "object": "response.input_tokens", - "input_tokens": _MANTLE_COUNT, - "litellm_count_source": "provider", - }, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_mantle_body(system=[{"type": "text", "text": "Be terse"}]),) - - -@pytest.mark.parametrize("case", LIFT_CASES.values(), ids=LIFT_CASES.keys()) -def test_messages_count_tokens_lifts_the_leading_system_run_for_mantle( - counting_proxy: OwnedProxy, mantle_port: int, case: LiftCase -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, count_request(model, case)) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (expected_count_body(_OPUS_BASE, case),) - - -def test_anthropic_sdk_count_tokens_lifts_the_leading_system_through_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - counted: Final = _anthropic_client(gateway).messages.count_tokens(model=model, messages=_LIFT_SDK_MESSAGES) - assert counted.input_tokens == _MANTLE_COUNT, counted - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) - - -def test_async_anthropic_sdk_count_tokens_lifts_the_leading_system_through_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - counted: Final = asyncio.run( - _async_anthropic_client(gateway).messages.count_tokens(model=model, messages=_LIFT_SDK_MESSAGES) - ) - assert counted.input_tokens == _MANTLE_COUNT, counted - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) - - -def test_utils_token_counter_call_endpoint_counts_a_leading_system_through_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request( - "POST", "/utils/token_counter", count_request(model, STRING_CASE), params={"call_endpoint": "true"} - ) - payload: Final = _payload(response) - assert (payload["total_tokens"], payload["tokenizer_type"]) == (_MANTLE_COUNT, "bedrock_mantle_api") - assert payload["original_response"] == {"input_tokens": _MANTLE_COUNT}, response.text - assert (payload["request_model"], payload["model_used"]) == (model, _OPUS), response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) - - -def test_openai_sdk_input_tokens_lifts_instructions_through_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - counted: Final = openai_client(gateway).responses.input_tokens.count( - model=model, input=USER_TEXT, instructions=INSTRUCTION - ) - assert counted.input_tokens == _MANTLE_COUNT, counted - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) - - -def test_async_openai_sdk_input_tokens_lifts_instructions_through_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - counted: Final = asyncio.run( - async_openai_client(gateway).responses.input_tokens.count( - model=model, input=USER_TEXT, instructions=INSTRUCTION - ) - ) - assert counted.input_tokens == _MANTLE_COUNT, counted - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) - - -def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_item_through_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request( - "POST", - "/v1/responses/input_tokens", - {"model": model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION}, - ) - assert _payload(response) == { - "object": "response.input_tokens", - "input_tokens": _MANTLE_COUNT, - "litellm_count_source": "provider", - }, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED, {"type": "text", "text": REMINDER}]),) - - -def test_messages_count_tokens_forwards_tools_beside_the_lifted_system_to_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, {**count_request(model, STRING_CASE), "tools": _TOOLS}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED], tools=_TOOLS),) - - -def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place_for_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - body: Final[dict[str, JsonValue]] = { - "model": model, - "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], - } - response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == ( - _mantle_body(messages=[USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], system=[LIFTED]), - ) - - -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 - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - body: Final[dict[str, JsonValue]] = {"model": model, "messages": [LEADING]} - response: Final = _count(gateway, body) - 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_surfaces_the_mantle_error_beside_a_non_text_system( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, {**count_request(model, STRING_CASE), "system": 5}) - 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),) - - -def test_messages_count_tokens_repeated_request_lifts_the_leading_system_each_time_for_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - answers: Final = tuple(_payload(_count(gateway, count_request(model, STRING_CASE))) for _ in range(2)) - assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2 - assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 - assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * 2 - - -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 - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - body: Final[dict[str, JsonValue]] = {"model": model, "messages": [{"role": "system"}, USER]} - response: Final = _count(gateway, body) - 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, "litellm_count_source": "provider"}, ( - follow_up.text - ) - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) - - -def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value_for_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - first: Final = json.dumps([USER]) - last: Final = json.dumps([LEADING, USER]) - response: Final = gateway.client.post( - "/v1/messages/count_tokens", - content=f'{{"model": "{model}", "messages": {first}, "messages": {last}}}', - headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, - ) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) - - -def test_disabled_token_counter_counts_a_leading_system_through_mantle( - gateway: Gateway, mantle_port: int, tmp_path: Path -) -> None: - with ExitStack() as stack: - runtime: Final = stack.enter_context(wire_server(_runtime)) - config: Final = _owned_config( - tmp_path / "disabled-token-counter-lift.yaml", runtime.url, {"disable_token_counter": True} - ) - owned: Final = stack.enter_context( - owned_proxy_process( - gateway, - tmp_path, - _mantle_environment(mantle_port), - config=config, - workers=2, - remove_environment=_INHERITED_BEARER, - ) - ) - with wire_server(_mantle(_strict), port=mantle_port) as strict: - counted: Final = _count(owned.gateway, count_request(_OWNED_OPUS, STRING_CASE)) - assert _payload(counted) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, ( - counted.text - ) - assert _mantle_bodies(strict) == (_MANTLE_LIFTED,) - assert len(_runtime_count_targets(runtime)) == 1 - - -def test_mantle_outage_between_concurrent_lifted_waves_errors_then_recovers( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ExitStack() as stack: - clients: Final = _clients(stack, str(gateway.client.base_url), 8) - pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) - runtime: Final = stack.enter_context(wire_server(_runtime)) - scenario: Final = stack.enter_context(gateway.scenario()) - model: Final = _deployment(scenario, runtime.url) - body: Final = count_request(model, STRING_CASE) - - def count(client: httpx.Client) -> tuple[int, JsonValue]: - return _counted_on(client, gateway.key, body) - - 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) - 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) - assert len(_runtime_count_targets(runtime)) == 3 * len(clients) - - -@pytest.mark.parametrize("status", [400, 403, 404, 500, 503]) -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 - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_rejecting(status)), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, _full(model)) - 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,) - - -@pytest.mark.parametrize( - "body", - [b'{"inputTokens": 7}', b"not json at all", b"{}"], - ids=["runtime_key", "not_json", "empty_object"], -) -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 - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(lambda _request: Reply(body=body)), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, _bare(model)) - 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_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 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_errors_when_mantle_never_answers( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - held: Final[SimpleQueue[str]] = SimpleQueue() - release: Final = threading.Event() - - def hold(request: Request) -> Reply: - held.put(request.target) - release.wait(timeout=120) - return Reply(drop_connection=True) - - with ExitStack() as stack: - (client,) = _clients(stack, str(gateway.client.base_url), 1, timeout=120) - pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=1)) - runtime: Final = stack.enter_context(wire_server(_runtime)) - mantle: Final = stack.enter_context(wire_server(_mantle(hold), port=mantle_port)) - stack.callback(release.set) - scenario: Final = stack.enter_context(gateway.scenario()) - model: Final = _deployment(scenario, runtime.url) - body: Final = _full(model) - 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) > 0 - assert not future.done() - 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_surfaces_the_mantle_error_for_a_non_text_system( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(_strict), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, {**_bare(model), "system": 5}) - 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),) - - -@pytest.mark.parametrize( - "tools", - [5, "", "x" * 5120, ["get_weather"]], - ids=["int", "empty_string", "5kb_string", "list_of_strings"], -) -def test_messages_count_tokens_rejects_malformed_tools_without_calling_either_peer( - counting_proxy: OwnedProxy, mantle_port: int, tools: JsonValue -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - refused: Final = _count(gateway, {**_bare(model), "tools": tools}) - assert 400 <= refused.status_code < 600, refused.text - assert "input_tokens" not in refused.text, refused.text - assert runtime.drain() == () and mantle.drain() == () - assert _generated_then_counted(gateway.client, gateway.key, model, _bare(model)) == (200, 200, _MANTLE_COUNT) - assert _mantle_bodies(mantle) == (_MANTLE_BARE,) - - -@pytest.mark.parametrize("fields", [{}, {"messages": []}], ids=["missing", "empty"]) -def test_messages_count_tokens_without_messages_is_rejected( - counting_proxy: OwnedProxy, mantle_port: int, fields: dict[str, JsonValue] -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, {"model": model, "system": _SYSTEM, "tools": _TOOLS, **fields}) - assert response.status_code == 400, response.text - assert "messages parameter is required" in response.text, response.text - assert runtime.drain() == () and mantle.drain() == () - - -def test_messages_count_tokens_unauthenticated_request_never_reaches_either_peer( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request( - "POST", "/v1/messages/count_tokens", _full(model), key="sk-not-a-key-this-proxy-issued" - ) - assert response.status_code == 401, response.text - assert "input_tokens" not in response.text, response.text - assert runtime.drain() == () and mantle.drain() == () - - -def test_messages_count_tokens_forwards_a_5kb_system_verbatim_to_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - system: Final = "Answer in one sentence. " * 214 - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, {**_bare(model), "system": system}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_mantle_body(system=system),) - - -def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value_each_to_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - fields: Final = f'"system": {json.dumps(_SYSTEM)}, "tools": {json.dumps(_TOOLS)}' - response: Final = gateway.client.post( - "/v1/messages/count_tokens", - content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}', - headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, - ) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text - assert len(_runtime_count_targets(runtime)) == 1 - (sent,) = mantle.drain() - assert _mantle_requests((sent,)) == (_MANTLE_FULL,) - assert (sent.body.count(b'"system"'), sent.body.count(b'"tools"')) == (1, 1), sent.body - - -def test_messages_count_tokens_leaves_empty_tools_and_system_out_of_the_mantle_body( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = _count(gateway, {**_bare(model), "tools": [], "system": ""}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text - assert len(_runtime_count_targets(runtime)) == 1 - assert _mantle_bodies(mantle) == (_MANTLE_BARE,) - - -def test_messages_count_tokens_repeated_request_reaches_both_peers_each_time( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2)) - assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2 - assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 - assert _mantle_bodies(mantle) == (_MANTLE_BARE,) * 2 - - -def test_disabled_token_counter_surfaces_the_mantle_error_instead_of_counting_locally( - gateway: Gateway, mantle_port: int, tmp_path: Path -) -> None: - with ExitStack() as stack: - runtime: Final = stack.enter_context(wire_server(_runtime)) - config: Final = _owned_config( - tmp_path / "disabled-token-counter.yaml", runtime.url, {"disable_token_counter": True} - ) - owned: Final = stack.enter_context( - owned_proxy_process( - gateway, - tmp_path, - _mantle_environment(mantle_port), - config=config, - workers=2, - remove_environment=_INHERITED_BEARER, - ) - ) - with wire_server(_mantle(_rejecting(403)), port=mantle_port) as refusing: - refused: Final = _count(owned.gateway, _full(_OWNED_OPUS)) - assert refused.status_code == 403, refused.text - assert _REJECTION in refused.text and "input_tokens" not in refused.text, refused.text - assert _mantle_bodies(refusing) == (_MANTLE_FULL,) - with wire_server(_mantle(), port=mantle_port) as counting: - assert _payload(_count(owned.gateway, _full(_OWNED_OPUS))) == { - "input_tokens": _MANTLE_COUNT, - "litellm_count_source": "provider", - } - assert _mantle_bodies(counting) == (_MANTLE_FULL,) - assert _payload(_count(owned.gateway, _full(_OWNED_SONNET))) == { - "input_tokens": _RUNTIME_COUNT, - "litellm_count_source": "provider", - } - rejected: Final = _count(owned.gateway, _full(_OWNED_NOVA)) - assert rejected.status_code == 400, rejected.text - assert _UNSUPPORTED in rejected.text and "input_tokens" not in rejected.text, rejected.text - assert counting.drain() == () - assert len(_runtime_count_targets(runtime)) == 4 - - -def test_mantle_outage_between_concurrent_waves_errors_then_recovers( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - with ExitStack() as stack: - clients: Final = _clients(stack, str(gateway.client.base_url), 8) - pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) - runtime: Final = stack.enter_context(wire_server(_runtime)) - scenario: Final = stack.enter_context(gateway.scenario()) - model: Final = _deployment(scenario, runtime.url) - body: Final = _full(model) - - def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]: - return _generated_then_counted(client, gateway.key, model, body) - - def count_only(client: httpx.Client) -> tuple[int, JsonValue]: - return _counted_on(client, gateway.key, body) - - with wire_server(_mantle(), port=mantle_port) as mantle: - 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)) - 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) - assert len(_runtime_count_targets(runtime)) == 3 * len(clients) - - -def test_slow_mantle_holds_concurrent_counts_without_stalling_the_proxy( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - held: Final[SimpleQueue[str]] = SimpleQueue() - release: Final = threading.Event() - with ExitStack() as stack: - clients: Final = _clients(stack, str(gateway.client.base_url), 6) - runtime: Final = stack.enter_context(wire_server(_runtime)) - mantle: Final = stack.enter_context(wire_server(_mantle(_holding(held, release, 20)), port=mantle_port)) - scenario: Final = stack.enter_context(gateway.scenario()) - pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) - stack.callback(release.set) - model: Final = _deployment(scenario, runtime.url) - futures: Final = tuple( - pool.submit(_generated_then_counted, client, gateway.key, model, _full(model)) for client in clients - ) - eventually(held.qsize, lambda size: size == len(clients), seconds=30) - assert gateway.request("GET", "/health/liveliness").status_code == 200 - assert _local_count(gateway, _full(model)) > 0 - assert not any(future.done() for future in futures) - release.set() - assert tuple(future.result(timeout=30) for future in futures) == ((200, 200, _MANTLE_COUNT),) * len(clients) - assert _mantle_bodies(mantle) == (_MANTLE_FULL,) * len(clients) - assert len(_runtime_count_targets(runtime)) == len(clients) - - -def test_worker_sigkill_mid_burst_leaves_the_sibling_counting( - gateway: Gateway, mantle_port: int, tmp_path: Path -) -> None: - held: Final[SimpleQueue[str]] = SimpleQueue() - release: Final = threading.Event() - with ExitStack() as stack: - runtime: Final = stack.enter_context(wire_server(_runtime)) - mantle: Final = stack.enter_context(wire_server(_mantle(_holding(held, release, 60)), port=mantle_port)) - config: Final = _owned_config(tmp_path / "worker-kill.yaml", runtime.url, {}) - owned: Final = stack.enter_context( - owned_proxy_process( - gateway, - tmp_path, - _mantle_environment(mantle_port), - config=config, - workers=2, - remove_environment=_INHERITED_BEARER, - ) - ) - body: Final = _full(_OWNED_OPUS) - proxy_url: Final = owned.gateway.client.base_url - workers: Final = eventually( - lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), - lambda pids: len(pids) == 2, - seconds=30, - ) - clients: Final = _clients(stack, str(proxy_url), 12) - pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) - stack.callback(release.set) - probed: Final[SimpleQueue[int]] = SimpleQueue() - futures: Final[tuple[Future[tuple[int, tuple[int, JsonValue] | None]], ...]] = tuple( - pool.submit(_probed_then_counted_or_dropped, client, owned.gateway.key, body, probed) for client in clients - ) - eventually(held.qsize, lambda size: size == len(clients), seconds=30) - ports: Final = frozenset(probed.get_nowait() for _ in clients) - shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & ports for pid in workers} - assert sum(map(len, shares.values())) == len(clients), shares - victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) - psutil.Process(victim).send_signal(signal.SIGKILL) - release.set() - for port, result in (future.result(timeout=60) for future in futures): - assert result == (None if port in shares[victim] else (200, _MANTLE_COUNT)), (port, result, shares) - second_wave: Final = _clients(stack, str(proxy_url), 6) - assert tuple(_counted_on(client, owned.gateway.key, body) for client in second_wave) == ( - (200, _MANTLE_COUNT), - ) * len(second_wave) - assert len(_mantle_bodies(mantle)) == len(clients) + len(second_wave) - eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120) - assert owned.process.poll() is None - - -def _streamed_text(text: str) -> str: - events: Final = tuple( - _JSON_OBJECT.validate_json(line.removeprefix("data: ")) - for line in text.splitlines() - if line.startswith("data: {") - ) - return "".join(_delta_content(event) for event in events) - - -def _delta_content(event: dict[str, JsonValue]) -> str: - choices: Final = event.get("choices") - if not isinstance(choices, list) or not choices: - return "" - delta: Final = object_value(choices[0]).get("delta") - if not isinstance(delta, dict): - return "" - content: Final = delta.get("content") - return content if isinstance(content, str) else "" - - -@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) -def test_chat_completions_on_the_same_deployment_still_generate( - counting_proxy: OwnedProxy, mantle_port: int, stream: bool -) -> None: - gateway: Final = counting_proxy.gateway - marker: Final = uuid.uuid4().hex - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request( - "POST", - "/v1/chat/completions", - { - "model": model, - "messages": [{"role": "user", "content": f"chat control marker-{marker}"}], - "stream": stream, - "cache": {"no-cache": True}, - }, - ) - assert response.status_code == 200, response.text - generated: Final = _streamed_text(response.text) if stream else response.text - assert answer(marker) in generated, response.text - assert not stream or response.text.rstrip().endswith("data: [DONE]"), response.text - (sent,) = runtime.drain() - assert target_of(sent) == f"/model/{_OPUS}/{'converse-stream' if stream else 'converse'}", sent.target - assert mantle.drain() == () - - -def test_messages_endpoint_on_the_same_deployment_still_generates(counting_proxy: OwnedProxy, mantle_port: int) -> None: - gateway: Final = counting_proxy.gateway - marker: Final = uuid.uuid4().hex - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request( - "POST", - "/v1/messages", - {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": f"marker-{marker}"}]}, - ) - assert response.status_code == 200, response.text - assert answer(marker) in response.text, response.text - (sent,) = runtime.drain() - assert target_of(sent) == f"/model/{_OPUS}/invoke", sent.target - assert mantle.drain() == () - - -def test_responses_endpoint_on_the_same_deployment_still_generates( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: - gateway: Final = counting_proxy.gateway - marker: Final = uuid.uuid4().hex - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": f"marker-{marker}"}) - assert response.status_code == 200, response.text - assert answer(marker) in response.text, response.text - (sent,) = runtime.drain() - assert target_of(sent) == f"/model/{_OPUS}/converse", sent.target - assert mantle.drain() == () - - -def test_utils_token_counter_local_mode_never_calls_either_peer(counting_proxy: OwnedProxy, mantle_port: int) -> None: - gateway: Final = counting_proxy.gateway - with ( - wire_server(_runtime) as runtime, - wire_server(_mantle(), port=mantle_port) as mantle, - gateway.scenario() as scenario, - ): - model: Final = _deployment(scenario, runtime.url) - assert _local_count(gateway, _full(model)) > 0 - assert runtime.drain() == () and mantle.drain() == () diff --git a/tests/integration/providers/test_chat_template_tool_count_wire.py b/tests/integration/providers/test_chat_template_tool_count_wire.py deleted file mode 100644 index d2b85e7168f..00000000000 --- a/tests/integration/providers/test_chat_template_tool_count_wire.py +++ /dev/null @@ -1,383 +0,0 @@ -import json -import re -import uuid -from collections.abc import Iterator, Mapping, Sequence -from contextlib import ExitStack -from dataclasses import dataclass -from itertools import chain -from pathlib import Path -from typing import Final - -import httpx -import pytest -import yaml -from integration._support.client import ( - JSON_OBJECT, - Gateway, - Scenario, - gateway_from_environment, - object_value, - string_value, -) -from integration._support.openai_wire import chat_reply -from integration._support.process import graceful_stop_seconds, owned_proxy_process -from integration._support.wire import Reply, Request, Wire, wire_server -from litellm.litellm_core_utils.default_encoding import encoding as cl100k_base -from pydantic import JsonValue - -pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) - -_API_KEY: Final = "scripted-provider-key" -_DEEPSEEK_V4: Final = "azure_ai/DeepSeek-V4-Flash-0731" -_QWEN3_5: Final = "databricks/databricks-qwen35-122b-a10b" -_GROK: Final = "azure_ai/grok-4-1" -_DEEPSEEK_V4_TOOL_PREAMBLE: Final = 222 -_QWEN3_5_TOOL_PREAMBLE: Final = 196 -_QWEN3_5_OPENED_SYSTEM_TURN: Final = 5 -_OPENAI_TOOL_OVERHEAD: Final = 9 -_SCHEMA: Final[dict[str, JsonValue]] = { - "type": "object", - "properties": {"query": {"type": "string"}}, - "required": ["query"], -} -_FUNCTION: Final[dict[str, JsonValue]] = {"name": "lookup", "parameters": _SCHEMA} -_OPENAI_NAMESPACE: Final = "namespace functions {\n\ntype lookup = (_: {\nquery: string,\n}) => any;\n\n} // namespace functions" -_ANTHROPIC_TOOL: Final[dict[str, JsonValue]] = {"name": "lookup", "input_schema": _SCHEMA} -_OPENAI_TOOL: Final[dict[str, JsonValue]] = {"type": "function", "function": _FUNCTION} -_RESPONSES_TOOL: Final[dict[str, JsonValue]] = {"type": "function", **_FUNCTION} -_NAMED_TOOL_CHOICE: Final[dict[str, JsonValue]] = {"type": "function", "function": {"name": "lookup"}} -_WEB_SEARCH_TOOL: Final[dict[str, JsonValue]] = {"type": "web_search_20250305", "name": "web_search", "max_uses": 3} -_PROMPT: Final = "Find record 14." -_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": _PROMPT}] -_INPUT_PRICE: Final = 0.001 -_OUTPUT_PRICE: Final = 0.002 -_TINY_BUDGET: Final = 0.000000001 -_ESTIMATED_COST: Final = re.compile(r"Estimated request cost: ([^\s,]+)") - - -def _tokens(text: str) -> int: - return len(cl100k_base.encode(text)) - - -def _deepseek_v4_tool_tokens(function: Mapping[str, JsonValue]) -> int: - return _DEEPSEEK_V4_TOOL_PREAMBLE + _tokens(json.dumps(function, ensure_ascii=False)) - - -def _qwen3_5_tool_tokens(function: Mapping[str, JsonValue], *, with_system: bool) -> int: - opened_system_turn: Final = 0 if with_system else _QWEN3_5_OPENED_SYSTEM_TURN - schema: Final = json.dumps({"type": "function", "function": function}, ensure_ascii=False) - return _QWEN3_5_TOOL_PREAMBLE + opened_system_turn + _tokens("\n" + schema) - - -def _completion(_request: Request) -> Reply: - return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", "scripted-model", "OK", stream=False) - - -def _deployment(scenario: Scenario, group: str, model: str, api_base: str) -> str: - name: Final = f"{group}-{uuid.uuid4().hex}" - created: Final = scenario.gateway.post( - "/model/new", - { - "model_name": name, - "litellm_params": {"model": model, "api_base": api_base, "api_key": _API_KEY}, - "model_info": {}, - }, - ) - scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) - return name - - -@dataclass(frozen=True, slots=True) -class _CountRoute: - path: str - field: str - tool: dict[str, JsonValue] - - def body(self, model: str, tools: Sequence[JsonValue]) -> dict[str, JsonValue]: - prompt: Final[dict[str, JsonValue]] = ( - {"input": _PROMPT} if self.path == "/v1/responses/input_tokens" else {"messages": _MESSAGES} - ) - return {"model": model, **prompt, **({"tools": list(tools)} if tools else {})} - - -_COUNT_ROUTES: Final[Mapping[str, _CountRoute]] = { - "messages-count-tokens": _CountRoute("/v1/messages/count_tokens", "input_tokens", _ANTHROPIC_TOOL), - "utils-token-counter": _CountRoute("/utils/token_counter", "total_tokens", _OPENAI_TOOL), - "responses-input-tokens": _CountRoute("/v1/responses/input_tokens", "input_tokens", _RESPONSES_TOOL), -} -_MESSAGES_COUNT: Final = _COUNT_ROUTES["messages-count-tokens"] - - -def _count(gateway: Gateway, route: _CountRoute, body: Mapping[str, JsonValue]) -> int: - response: Final = gateway.request("POST", route.path, body) - assert response.status_code == 200, response.text - counted: Final = JSON_OBJECT.validate_json(response.content)[route.field] - assert isinstance(counted, int), response.text - return counted - - -@pytest.mark.parametrize("route", sorted(_COUNT_ROUTES)) -def test_count_routes_count_deepseek_v4_tools_the_way_its_chat_template_renders_them( - gateway: Gateway, route: str -) -> None: - selected: Final = _COUNT_ROUTES[route] - with wire_server(_completion) as wire, gateway.scenario() as scenario: - model: Final = _deployment(scenario, "count", _DEEPSEEK_V4, wire.url) - without_tools: Final = _count(gateway, selected, selected.body(model, ())) - with_tool: Final = _count(gateway, selected, selected.body(model, (selected.tool,))) - upstream: Final = wire.drain() - assert upstream == () - assert with_tool - without_tools == _deepseek_v4_tool_tokens(_FUNCTION) - - -@pytest.mark.parametrize("with_system", (False, True), ids=("no-system", "system")) -def test_messages_count_tokens_counts_qwen3_5_tools_the_way_its_chat_template_renders_them( - gateway: Gateway, with_system: bool -) -> None: - system: Final[dict[str, JsonValue]] = {"system": "Be brief."} if with_system else {} - with wire_server(_completion) as wire, gateway.scenario() as scenario: - model: Final = _deployment(scenario, "count", _QWEN3_5, wire.url) - without_tools: Final = _count(gateway, _MESSAGES_COUNT, {**_MESSAGES_COUNT.body(model, ()), **system}) - with_tool: Final = _count( - gateway, _MESSAGES_COUNT, {**_MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,)), **system} - ) - upstream: Final = wire.drain() - assert upstream == () - assert with_tool - without_tools == _qwen3_5_tool_tokens(_FUNCTION, with_system=with_system) - - -def test_messages_count_tokens_skips_web_search_and_counts_unmapped_tool_keys_for_deepseek_v4( - gateway: Gateway, -) -> None: - deferred_function: Final[dict[str, JsonValue]] = {"name": "lookup", "parameters": {**_SCHEMA, "defer_loading": True}} - with wire_server(_completion) as wire, gateway.scenario() as scenario: - model: Final = _deployment(scenario, "count", _DEEPSEEK_V4, wire.url) - plain: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,))) - with_web_search: Final = _count( - gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL, _WEB_SEARCH_TOOL)) - ) - deferred: Final = _count( - gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ({**_ANTHROPIC_TOOL, "defer_loading": True},)) - ) - upstream: Final = wire.drain() - assert upstream == () - assert with_web_search == plain - assert deferred - plain == _deepseek_v4_tool_tokens(deferred_function) - _deepseek_v4_tool_tokens(_FUNCTION) - - -def test_messages_count_tokens_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None: - with wire_server(_completion) as wire, gateway.scenario() as scenario: - model: Final = _deployment(scenario, "count", _GROK, wire.url) - without_tools: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ())) - plain: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,))) - deferred: Final = _count( - gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ({**_ANTHROPIC_TOOL, "defer_loading": True},)) - ) - upstream: Final = wire.drain() - assert upstream == () - assert plain - without_tools == _OPENAI_TOOL_OVERHEAD + _tokens(_OPENAI_NAMESPACE) - assert deferred == plain - - -@dataclass(frozen=True, slots=True) -class _Endpoint: - path: str - body: dict[str, JsonValue] - tools: dict[str, JsonValue] - - -_ENDPOINTS: Final[Mapping[str, _Endpoint]] = { - "chat-completions": _Endpoint( - "/v1/chat/completions", - {"messages": _MESSAGES, "max_tokens": 1}, - {"tools": [_OPENAI_TOOL], "tool_choice": _NAMED_TOOL_CHOICE}, - ), - "messages": _Endpoint( - "/v1/messages", - {"messages": _MESSAGES, "max_tokens": 1}, - {"tools": [_ANTHROPIC_TOOL], "tool_choice": {"type": "auto"}}, - ), - "responses": _Endpoint( - "/v1/responses", - {"input": _PROMPT, "max_output_tokens": 16}, - {"tools": [_RESPONSES_TOOL]}, - ), -} - - -def _send(gateway: Gateway, endpoint: _Endpoint, group: str, key: str, *, tools: bool) -> httpx.Response: - return gateway.request( - "POST", - endpoint.path, - {"model": group, **endpoint.body, **(endpoint.tools if tools else {})}, - key=key, - ) - - -@dataclass(frozen=True, slots=True) -class _BudgetRig: - gateway: Gateway - wire: Wire - group: str - - -def _strict_budget_config(directory: Path, group: str, api_base: str) -> Path: - base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) - config: Final = { - **base, - "model_list": [ - { - "model_name": group, - "litellm_params": { - "model": _DEEPSEEK_V4, - "api_base": api_base, - "api_key": _API_KEY, - "input_cost_per_token": _INPUT_PRICE, - "output_cost_per_token": _OUTPUT_PRICE, - }, - } - ], - "general_settings": {**object_value(base["general_settings"]), "fail_closed_budget_enforcement": True}, - } - path: Final = directory / f"chat-template-tool-budget-{uuid.uuid4().hex[:8]}.yaml" - path.write_text(yaml.safe_dump(config)) - return path - - -@pytest.fixture(scope="module") -def strict_budget(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_BudgetRig]: - directory: Final = tmp_path_factory.mktemp("chat-template-tool-budget") - group: Final = f"deepseek-v4-flash-{uuid.uuid4().hex}" - with ExitStack() as stack: - rig_gateway: Final = stack.enter_context(gateway_from_environment()) - wire: Final = stack.enter_context(wire_server(_completion)) - owned: Final = stack.enter_context( - owned_proxy_process(rig_gateway, directory, {}, config=_strict_budget_config(directory, group, wire.url)) - ) - yield _BudgetRig(owned.gateway, wire, group) - - -def _estimated_cost(response: httpx.Response) -> float: - assert response.status_code in (400, 422, 429), f"{response.status_code} {response.text}" - error: Final = object_value(JSON_OBJECT.validate_json(response.content)["error"]) - estimate: Final = _ESTIMATED_COST.search(string_value(error["message"])) - assert estimate is not None, response.text - return float(estimate[1]) - - -@pytest.mark.parametrize("endpoint", ("chat-completions", "messages")) -def test_budget_reservation_estimates_deepseek_v4_tools_the_way_its_chat_template_renders_them( - strict_budget: _BudgetRig, endpoint: str -) -> None: - selected: Final = _ENDPOINTS[endpoint] - with strict_budget.gateway.scenario() as scenario: - key: Final = scenario.key(models=[strict_budget.group], max_budget=_TINY_BUDGET) - without_tools: Final = _estimated_cost(_send(strict_budget.gateway, selected, strict_budget.group, key, tools=False)) - with_tools: Final = _estimated_cost(_send(strict_budget.gateway, selected, strict_budget.group, key, tools=True)) - assert strict_budget.wire.drain() == () - assert with_tools - without_tools == pytest.approx(_deepseek_v4_tool_tokens(_FUNCTION) * _INPUT_PRICE) - - -def _project_key(scenario: Scenario, group: str) -> str: - team: Final = scenario.team() - project: Final = scenario.project(team, model_itpm_limit={group: _deepseek_v4_tool_tokens(_FUNCTION)}) - return scenario.key(team_id=team, project_id=project) - - -@pytest.mark.parametrize("endpoint", sorted(_ENDPOINTS)) -def test_project_itpm_reserves_deepseek_v4_tools_the_way_its_chat_template_renders_them( - gateway: Gateway, endpoint: str -) -> None: - with wire_server(_completion) as wire, gateway.scenario() as scenario: - group: Final = _deployment(scenario, "deepseek-v4-flash", _DEEPSEEK_V4, wire.url) - response: Final = _send(gateway, _ENDPOINTS[endpoint], group, _project_key(scenario, group), tools=True) - upstream: Final = wire.drain() - assert response.status_code == 429, response.text - assert "model_per_project_itpm" in response.text, response.text - assert upstream == () - - -def test_project_itpm_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None: - chat: Final = _ENDPOINTS["chat-completions"] - with wire_server(_completion) as wire, gateway.scenario() as scenario: - group: Final = _deployment(scenario, "grok-4-1", _GROK, wire.url) - response: Final = gateway.request( - "POST", - chat.path, - {"model": group, **chat.body, "tools": [_OPENAI_TOOL]}, - key=_project_key(scenario, group), - ) - upstream: Final = wire.drain() - assert response.status_code == 200, response.text - assert len(upstream) == 1, upstream - - -_FILES: Final = 3 -_READ_FILE_FUNCTION: Final[dict[str, JsonValue]] = { - "name": "read_file", - "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}, -} -_READ_FILE_TOOL: Final[dict[str, JsonValue]] = { - "name": "read_file", - "input_schema": _READ_FILE_FUNCTION["parameters"], -} - - -def _read_file_round(marker: str, index: int) -> tuple[JsonValue, ...]: - return ( - {"role": "user", "content": f"Read file {index}."}, - { - "role": "assistant", - "content": [ - {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}.txt"}} - ], - }, - { - "role": "user", - "content": [ - { - "type": "tool_result", - "tool_use_id": f"toolu_{index}", - "content": f"short file body.", - } - ], - }, - ) - - -def _agent_body(group: str, marker: str) -> dict[str, JsonValue]: - rounds: Final = tuple(chain.from_iterable(_read_file_round(marker, index) for index in range(_FILES))) - return { - "model": group, - "max_tokens": 16, - "tools": [_READ_FILE_TOOL], - "messages": [*rounds, {"role": "user", "content": "Summarize what you read."}], - "context_management": { - "edits": [ - { - "type": "clear_tool_uses_20250919", - "trigger": {"type": "input_tokens", "value": _deepseek_v4_tool_tokens(_READ_FILE_FUNCTION)}, - "keep": {"type": "tool_uses", "value": 1}, - } - ] - }, - } - - -@pytest.mark.parametrize( - ("model", "files_kept"), - ((_DEEPSEEK_V4, (False, False, True)), (_GROK, (True, True, True))), - ids=("deepseek-v4", "grok-4-1"), -) -def test_messages_context_management_trigger_counts_tools_the_way_the_models_chat_template_renders_them( - gateway: Gateway, model: str, files_kept: tuple[bool, ...] -) -> None: - marker: Final = uuid.uuid4().hex - with wire_server(_completion) as wire, gateway.scenario() as scenario: - group: Final = _deployment(scenario, "context", model, wire.url) - response: Final = gateway.request("POST", "/v1/messages", _agent_body(group, marker)) - upstream: Final = wire.drain() - assert response.status_code == 200, response.text - assert len(upstream) == 1, upstream - sent: Final = upstream[0].body.decode() - assert tuple(f"" in sent for index in range(_FILES)) == files_kept diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index 16a5b8685ad..2b23ccdf4f2 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -391,28 +391,3 @@ def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: O ) assert response.status_code == 200, response.text assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English." - - -@pytest.mark.parametrize( - "error", - [ - pytest.param( - {"message": "Invalid input", "type": "invalid_request_error", "param": "input", "code": "invalid_value"}, - id="full", - ), - pytest.param({"message": "Invalid input", "type": "invalid_request_error"}, id="bare"), - ], -) -def test_responses_input_tokens_passes_an_openai_count_error_through( - counting_proxy: OwnedProxy, error: dict[str, JsonValue] -) -> None: - def rejecting(_request: Request) -> Reply: - return Reply(status=400, body=json.dumps({"error": error}).encode()) - - gateway: Final = counting_proxy.gateway - with wire_server(rejecting) as wire, gateway.scenario() as scenario: - model: Final = _deployment(scenario, wire.url) - response: Final = gateway.request("POST", "/v1/responses/input_tokens", {"model": model, "input": "hello"}) - assert response.status_code == 400, response.text - assert response.json() == {"error": {"param": None, "code": None, **error}}, response.text - assert _count_body(wire)["input"] == [{"role": "user", "content": "hello"}] diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 8de93256f6c..5fe82d37e9e 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -65,7 +65,7 @@ class _GeminiPeer: return requests -def _deployment(gateway: Gateway, scenario: Scenario, api_base: str, **parameters: JsonValue) -> str: +def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str: return scenario.model( model=f"vertex_ai/{_MODEL}", api_base=api_base, @@ -73,7 +73,6 @@ def _deployment(gateway: Gateway, scenario: Scenario, api_base: str, **parameter vertex_project=_PROJECT, vertex_location=_LOCATION, vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), - **parameters, ) @@ -221,8 +220,8 @@ _CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingPro ) -def _ai_studio_deployment(scenario: Scenario, api_base: str, **parameters: JsonValue) -> str: - return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key", **parameters) +def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key") def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: @@ -595,110 +594,3 @@ def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gatew tuple(_json_body(request) for request in requests) == ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9 ) - - -_CONTENTS: Final[JsonValue] = [{"role": "user", "parts": [{"text": _USER}]}] -_OPENAI_TOOLS: Final[JsonValue] = [ - { - "type": "function", - "function": {"name": "noop", "description": "Return no output", "parameters": {"type": "object"}}, - } -] - - -def _provider_deployment( - gateway: Gateway, scenario: Scenario, provider: str, api_base: str, **parameters: JsonValue -) -> str: - if provider == "vertex_ai": - return _deployment(gateway, scenario, api_base, **parameters) - return _ai_studio_deployment(scenario, api_base, **parameters) - - -def _count(gateway: Gateway, body: Mapping[str, JsonValue], call_endpoint: bool = True) -> httpx.Response: - return gateway.request( - "POST", "/utils/token_counter", body, params={"call_endpoint": "true" if call_endpoint else "false"} - ) - - -def _total_tokens(response: httpx.Response) -> JsonValue: - return _JSON_OBJECT.validate_json(response.content)["total_tokens"] - - -@pytest.mark.parametrize( - ("provider", "empty_conversation"), - (("vertex_ai", {}), ("gemini", {}), ("gemini", {"messages": []}), ("gemini", {"contents": []})), - ids=("vertex-prompt-only", "ai-studio-prompt-only", "ai-studio-empty-messages", "ai-studio-empty-contents"), -) -def test_utils_token_counter_sends_a_prompt_as_the_user_turn_of_its_messages_form( - gateway: Gateway, provider: str, empty_conversation: Mapping[str, JsonValue] -) -> None: - peer: Final = _GeminiPeer() - with wire_server(peer) as wire, gateway.scenario() as scenario: - model: Final = _provider_deployment(gateway, scenario, provider, wire.url) - prompted: Final = _count(gateway, {"model": model, "prompt": _USER, **empty_conversation}) - prompted_requests: Final = _count_requests(wire) - as_message: Final = _count(gateway, {"model": model, "messages": _MESSAGES}) - message_requests: Final = _count_requests(wire) - assert prompted.status_code == 200, prompted.text - assert as_message.status_code == 200, as_message.text - assert _total_tokens(prompted) == _COUNT, prompted.text - assert len(prompted_requests) == 1 - assert len(message_requests) == 1 - assert _counted_contents(_json_body(prompted_requests[0])) == _counted_contents(_json_body(message_requests[0])) - - -@pytest.mark.parametrize( - ("conversation", "call_endpoint"), - (({"messages": _MESSAGES}, True), ({"contents": _CONTENTS}, True), ({"messages": _MESSAGES}, False)), - ids=("provider-count-messages", "provider-count-contents", "local-count-messages"), -) -def test_utils_token_counter_rejects_a_prompt_next_to_a_conversation_without_calling_the_provider( - gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool -) -> None: - peer: Final = _GeminiPeer() - with wire_server(peer) as wire, gateway.scenario() as scenario: - model: Final = _ai_studio_deployment(scenario, wire.url) - response: Final = _count(gateway, {"model": model, "prompt": "hi", **conversation}, call_endpoint) - requests: Final = _count_requests(wire) - assert response.status_code == 400, response.text - assert _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_json(response.content)["error"])["param"] == "prompt" - assert requests == () - - -@pytest.mark.parametrize("call_endpoint", (True, False), ids=("provider-count", "local-count")) -def test_utils_token_counter_counts_only_the_conversation_when_drop_params_drops_the_prompt( - gateway: Gateway, call_endpoint: bool -) -> None: - peer: Final = _GeminiPeer() - with wire_server(peer) as wire, gateway.scenario() as scenario: - model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=True) - mixed: Final = _count(gateway, {"model": model, "prompt": "hi", "messages": _MESSAGES}, call_endpoint) - mixed_requests: Final = _count_requests(wire) - alone: Final = _count(gateway, {"model": model, "messages": _MESSAGES}, call_endpoint) - alone_requests: Final = _count_requests(wire) - assert mixed.status_code == 200, mixed.text - assert alone.status_code == 200, alone.text - assert _total_tokens(mixed) == _total_tokens(alone) - assert [_counted_contents(_json_body(request)) for request in mixed_requests] == [ - _counted_contents(_json_body(request)) for request in alone_requests - ] - - -@pytest.mark.parametrize( - "extra", - ({"system": "Follow the system instruction."}, {"tools": _OPENAI_TOOLS}), - ids=("system", "tools"), -) -def test_utils_token_counter_counts_a_local_prompt_with_system_or_tools_like_its_messages_form( - gateway: Gateway, extra: Mapping[str, JsonValue] -) -> None: - peer: Final = _GeminiPeer() - with wire_server(peer) as wire, gateway.scenario() as scenario: - model: Final = _ai_studio_deployment(scenario, wire.url) - prompted: Final = _count(gateway, {"model": model, "prompt": _USER, **extra}, call_endpoint=False) - as_message: Final = _count(gateway, {"model": model, "messages": _MESSAGES, **extra}, call_endpoint=False) - requests: Final = _count_requests(wire) - assert prompted.status_code == 200, prompted.text - assert as_message.status_code == 200, as_message.text - assert _total_tokens(prompted) == _total_tokens(as_message) - assert requests == () diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index d5c841d9335..73c135944e0 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -468,46 +468,10 @@ def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> Non assert sent["tools"] == _RESPONSES_TOOLS, sent -def test_responses_input_tokens_types_a_peer_error_by_its_status(gateway: Gateway) -> None: - with wire_server(_peer(_rejecting(404, "not_found_error"))) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire.url) - response: Final = _RESPONSES_ROUTE.send(gateway, model) - assert response.status_code == 404, response.text - assert response.json() == _RESPONSES_ROUTE.rejection(_REJECTION), response.text - assert _counted_bodies(wire) == (_PEER_BARE,) -@pytest.mark.parametrize( - ("backend", "route"), - [ - pytest.param(_GLM_BACKEND, _MESSAGES_ROUTE, id="glm-messages"), - pytest.param(_GLM_BACKEND, _RESPONSES_ROUTE, id="glm-responses"), - pytest.param(_GLM_BACKEND, _GOOGLE_ROUTE, id="glm-google"), - pytest.param(_GLM_BACKEND, _UTILS_ROUTE, id="glm-utils"), - pytest.param(_MISTRAL_BACKEND, _MESSAGES_ROUTE, id="mistral-messages"), - ], -) -def test_partner_without_a_count_route_answers_400_without_calling_vertex( - gateway: Gateway, backend: str, route: _CountRoute -) -> None: - with wire_server(_peer()) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{backend}") - response: Final = route.send(gateway, model) - assert response.status_code == 400, response.text - payload: Final = _JSON_OBJECT.validate_json(response.content) - message: Final = _error_message(payload) - assert backend in message, message - assert payload == route.rejection(message), response.text - assert wire.drain() == () -def test_llama_count_still_reaches_the_meta_count_route(gateway: Gateway) -> None: - with wire_server(_peer()) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{_LLAMA_BACKEND}") - response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text - (sent,) = _count_bodies(wire.drain(), _LLAMA_COUNT_TARGET) - assert sent["messages"] == _MESSAGES, sent @pytest.mark.timeout(2 * graceful_stop_seconds() + 120) diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index 8ed5237868a..58312faa1f9 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -1,4 +1,11 @@ model_list: + - model_name: anthropic/claude-opus-5-5 + litellm_params: + model: anthropic/claude-opus-5-5 + api_base: http://127.0.0.1:8191 + api_key: synthetic-anthropic-key + model_info: + token_count_framing: 0 - model_name: perplexity/pplx-decider-v1-27b litellm_params: model: perplexity/pplx-decider-v1-27b diff --git a/tests/test_litellm/proxy/google_endpoints/test_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_endpoints.py index 17af52b5579..92001118e2c 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_endpoints.py @@ -2,84 +2,14 @@ Test for google_endpoints/endpoints.py """ -from typing import Final - import pytest +import sys, os from dotenv import load_dotenv -from pydantic import JsonValue -from starlette.requests import Request -from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request, google_count_tokens + +from litellm.proxy.google_endpoints.endpoints import google_count_tokens from litellm.types.llms.vertex_ai import TokenCountDetailsResponse - - -@pytest.mark.parametrize( - ("system_key", "schema_key"), - (("systemInstruction", "parameters"), ("system_instruction", "parametersJsonSchema")), -) -def test_google_count_tokens_translates_system_and_function_declarations( - system_key: str, schema_key: str -) -> None: - schema: Final = {"type": "object", "properties": {"query": {"type": "string"}}} - request_body: Final[dict[str, JsonValue]] = { - "contents": [{"role": "user", "parts": [{"text": "Look this up"}]}], - system_key: {"parts": [{"text": "Use the provided tool."}]}, - "tools": [ - { - "functionDeclarations": [ - { - "name": "lookup", - "description": "Find a record", - schema_key: schema, - } - ] - } - ], - } - request: Final = _google_count_tokens_request("claude-sonnet", request_body) - - assert request.system == [{"type": "text", "text": "Use the provided tool."}] - assert request.tools == [ - { - "name": "lookup", - "description": "Find a record", - "input_schema": schema, - } - ] - assert request.messages == [{"role": "user", "content": "Look this up"}] - - -@pytest.mark.parametrize( - ("schema_key", "expected_schema"), - ( - ("parameters", {"type": "object", "properties": {"query": {"type": "string"}}}), - ("parametersJsonSchema", {"type": "OBJECT", "properties": {"query": {"type": "STRING"}}}), - ), -) -def test_google_count_tokens_normalizes_parameters_only( - schema_key: str, expected_schema: dict[str, JsonValue] -) -> None: - request: Final = _google_count_tokens_request( - "claude-sonnet", - { - "contents": [{"role": "user", "parts": [{"text": "Look this up"}]}], - "tools": [ - { - "functionDeclarations": [ - { - "name": "lookup", - schema_key: { - "type": "OBJECT", - "properties": {"query": {"type": "STRING"}}, - }, - } - ] - } - ], - }, - ) - - assert request.tools == [{"name": "lookup", "input_schema": expected_schema}] +from starlette.requests import Request load_dotenv() diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py index ee165d7b204..dc4e929637e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py @@ -383,7 +383,8 @@ def test_token_counter_fallback_counts_tools_system_and_anthropic_blocks(client, assert full > bare -def test_token_counter_fallback_counts_a_prompt_with_tools_like_its_messages_form(client, auth_as, monkeypatch): +def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, monkeypatch): + """Regression: a ``prompt`` request carrying ``tools`` but no ``messages`` still counts, because the fallback attaches tools only when counting messages (``token_counter`` rejects tools on the text path).""" monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False) prompt = "count the tokens in this sentence please" @@ -404,6 +405,4 @@ def test_token_counter_fallback_counts_a_prompt_with_tools_like_its_messages_for ) assert response.status_code == 200, response.text - assert response.json()["total_tokens"] == litellm.token_counter( - model="claude-fable-5", messages=[{"role": "user", "content": prompt}], tools=tools - ) + assert response.json()["total_tokens"] == litellm.token_counter(model="claude-fable-5", text=prompt) diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index d7648a08b6a..e684aa55b33 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -3,7 +3,6 @@ Test for response_api_endpoints/endpoints.py """ import unittest -from collections.abc import Mapping from typing import Any, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -2429,219 +2428,6 @@ class TestResponsesInputTokens: assert response.status_code == 429, response.text assert response.json()["error"]["message"] == "rate limited" - @pytest.mark.parametrize( - "provider_body, status_code, expected_error", - [ - pytest.param( - { - "error": { - "message": "Missing required parameter: 'tools[0].name'.", - "type": "invalid_request_error", - "param": "tools[0].name", - "code": "missing_required_parameter", - } - }, - 400, - { - "message": "Missing required parameter: 'tools[0].name'.", - "type": "invalid_request_error", - "param": "tools[0].name", - "code": "missing_required_parameter", - }, - id="openai_error_passes_through", - ), - pytest.param( - { - "error": { - "message": "Missing required parameter: 'tools[0].name'.", - "type": "invalid_request_error", - "param": "tools[0].name", - } - }, - 400, - { - "message": "Missing required parameter: 'tools[0].name'.", - "type": "invalid_request_error", - "param": "tools[0].name", - "code": None, - }, - id="openai_error_without_code_keeps_type_and_param", - ), - pytest.param( - { - "type": "error", - "error": {"type": "invalid_request_error", "message": "tools.0: Input tag 'function' is invalid"}, - "request_id": "req_vrtx_1", - }, - 400, - { - "message": "tools.0: Input tag 'function' is invalid", - "type": "invalid_request_error", - "param": None, - "code": None, - }, - id="anthropic_error_message_unwrapped", - ), - pytest.param( - { - "type": "error", - "error": {"type": "not_found_error", "message": "model: claude-typo"}, - "request_id": "req_vrtx_2", - }, - 404, - {"message": "model: claude-typo", "type": "invalid_request_error", "param": None, "code": None}, - id="anthropic_error_type_follows_the_status", - ), - pytest.param( - {"error": {"code": 404, "message": "Publisher model not found.", "status": "NOT_FOUND"}}, - 404, - {"message": "Publisher model not found.", "type": "invalid_request_error", "param": None, "code": None}, - id="google_error_message_unwrapped", - ), - pytest.param( - "litellm.APIError: Google Gen AI Studio API error: 404 - " - + json.dumps( - {"error": {"code": 404, "message": "models/gemini-typo is not found.", "status": "NOT_FOUND"}} - ), - 404, - { - "message": "models/gemini-typo is not found.", - "type": "invalid_request_error", - "param": None, - "code": None, - }, - id="prefixed_gemini_error_message_unwrapped", - ), - pytest.param( - "litellm.BadRequestError: AzureException - " - + json.dumps( - { - "error": { - "message": "Unrecognized request argument supplied: foo", - "type": "invalid_request_error", - "param": "foo", - "code": "unknown_parameter", - } - } - ), - 400, - { - "message": "Unrecognized request argument supplied: foo", - "type": "invalid_request_error", - "param": "foo", - "code": "unknown_parameter", - }, - id="prefixed_openai_error_passes_through", - ), - pytest.param( - {"message": "The provided model doesn't support counting tokens."}, - 400, - { - "message": "The provided model doesn't support counting tokens.", - "type": "invalid_request_error", - "param": None, - "code": None, - }, - id="bedrock_error_message_unwrapped", - ), - pytest.param( - "BedrockException - " + json.dumps({"message": "The provided model doesn't support counting tokens."}), - 400, - { - "message": "The provided model doesn't support counting tokens.", - "type": "invalid_request_error", - "param": None, - "code": None, - }, - id="prefixed_bedrock_error_message_unwrapped", - ), - pytest.param( - "tool schema {unclosed is invalid", - 400, - { - "message": "tool schema {unclosed is invalid", - "type": "invalid_request_error", - "param": None, - "code": None, - }, - id="plain_text_with_a_brace_kept_whole", - ), - pytest.param( - "upstream count service unavailable", - 503, - { - "message": "upstream count service unavailable", - "type": "internal_server_error", - "param": None, - "code": None, - }, - id="plain_text_error", - ), - ], - ) - def test_provider_count_error_answers_in_openai_error_shape( - self, - provider_body: str | Mapping[str, object], - status_code: int, - expected_error: Mapping[str, str | None], - ) -> None: - from litellm.proxy._types import ProxyException - - failing_counter: Final = AsyncMock( - side_effect=ProxyException( - message=provider_body if isinstance(provider_body, str) else json.dumps(provider_body), - type="token_counting_error", - param="model", - code=str(status_code), - ) - ) - response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) - - assert response.status_code == status_code, response.text - assert response.json() == {"error": expected_error} - - def test_non_provider_count_error_keeps_its_own_type(self) -> None: - from litellm.proxy._types import ProxyException - - failing_counter: Final = AsyncMock( - side_effect=ProxyException( - message="Token counting is disabled and no provider API result available", - type="token_counting_disabled", - param="model", - code=503, - ) - ) - response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) - - assert response.status_code == 503, response.text - assert response.json()["error"]["type"] == "token_counting_disabled" - - def test_provider_count_error_names_the_call_id_when_opted_in(self) -> None: - from litellm.proxy._types import ProxyException - - failing_counter: Final = AsyncMock( - side_effect=ProxyException( - message="LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas.", - type="token_counting_error", - param="model", - code="400", - headers={"x-litellm-call-id": "call-9564"}, - ) - ) - with patch("litellm.proxy.proxy_server.general_settings", {"include_call_id_in_error_body": True}): - response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) - - assert response.status_code == 400, response.text - assert response.json() == { - "error": { - "message": "LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas.", - "type": "invalid_request_error", - "param": None, - "code": None, - "litellm_call_id": "call-9564", - } - } - def test_responses_routes_document_response_models_in_openapi_schema(): from typing import cast diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py deleted file mode 100644 index 51249e01d4d..00000000000 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py +++ /dev/null @@ -1,31 +0,0 @@ -from typing import Final - -import pytest - -from litellm.litellm_core_utils.prompt_templates.chat_template_tools import chat_template_tool_section - -_DESCRIBED_FUNCTION: Final = {"name": "f", "description": "é", "parameters": {}} -_BARE_FUNCTION: Final = {"name": "g", "parameters": {}} - - -# Tool schemas as rendered by render_tools in encoding/encoding_dsv4.py (deepseek-ai/DeepSeek-V4-Flash) and by -# chat_template.jinja (Qwen/Qwen3.5-122B-A10B), both read on Hugging Face on 2026-10-17 -@pytest.mark.parametrize( - ("model", "rendered"), - [ - ( - "DeepSeek-V4-Flash-0731", - '{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}', - ), - ( - "databricks-qwen35-122b-a10b", - '\n{"type": "function", "function": {"name": "f", "description": "é", "parameters": {}}}' - '\n{"type": "function", "function": {"name": "g", "parameters": {}}}', - ), - ], -) -def test_chat_template_tool_section_renders_schemas_like_the_models_template(model: str, rendered: str): - section: Final = chat_template_tool_section(model) - - assert section is not None - assert section.render_tool_schemas((_DESCRIBED_FUNCTION, _BARE_FUNCTION)) == rendered diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py index 9e3704a9df0..d1e23a17747 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py @@ -11,7 +11,6 @@ import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( CONVERTED_SYSTEM_NOTE, - anthropic_system_blocks, place_mid_conversation_system, split_leading_system_run, ) @@ -391,23 +390,3 @@ def test_flagged_placement_keeps_a_system_before_an_assistant_turn_whose_empty_t ) assert _roles(placed) == ["user", "system", "assistant", "user"] - - -def test_anthropic_system_blocks_keeps_text_parts_with_their_cache_control_and_drops_the_rest(): - run = [ - {"role": "system", "content": "one", "cache_control": {"type": "ephemeral", "ttl": "1h"}}, - { - "role": "system", - "content": [ - {"type": "text", "text": "two", "cache_control": {"type": "ephemeral"}}, - {"type": "text", "text": ""}, - {"type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="}}, - ], - }, - {"role": "system", "content": ""}, - ] - - assert anthropic_system_blocks(run) == ( - {"type": "text", "text": "one", "cache_control": {"type": "ephemeral", "ttl": "1h"}}, - {"type": "text", "text": "two", "cache_control": {"type": "ephemeral"}}, - ) diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 8d2884756cb..1285a291650 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -9,9 +9,8 @@ import subprocess import sys import threading import time -from collections.abc import Callable, Mapping +from collections.abc import Mapping from concurrent.futures import Future, wait -from copy import deepcopy from itertools import accumulate, chain from pathlib import Path from typing import Final @@ -40,7 +39,6 @@ from litellm.litellm_core_utils.token_counter import ( offload_token_count, ) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new -from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter from tests.large_text import text from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, @@ -445,255 +443,6 @@ def test_token_counter_with_tools(message_count_pair): ), f"Expected {expected_tokens} tokens, got {counted_tokens}." -def test_token_counter_counts_gemini_function_declarations(): - openai_tools: Final = [ - { - "type": "function", - "function": { - "name": "lookup_weather", - "description": "Find current weather conditions for a location", - "parameters": { - "type": "object", - "properties": { - "location": {"type": "string", "description": "City and region"}, - "units": {"type": "string", "enum": ["celsius", "fahrenheit"]}, - }, - "required": ["location"], - }, - }, - } - ] - gemini_tools: Final = litellm.utils.get_optional_params( - model="gemini-2.5-pro", - custom_llm_provider="gemini", - tools=openai_tools, - )["tools"] - camel_case_tools: Final = [{"functionDeclarations": gemini_tools[0]["function_declarations"]}] - - openai_tokens: Final = token_counter_new( - model="gemini-2.5-pro", - messages=[{"role": "user", "content": "What's the weather?"}], - tools=openai_tools, - ) - gemini_tokens: Final = token_counter_new( - model="gemini-2.5-pro", - messages=[{"role": "user", "content": "What's the weather?"}], - tools=gemini_tools, - ) - camel_case_tokens: Final = token_counter_new( - model="gemini-2.5-pro", - messages=[{"role": "user", "content": "What's the weather?"}], - tools=camel_case_tools, - ) - - assert openai_tokens == gemini_tokens == camel_case_tokens - - -def test_token_counter_skips_non_mapping_tools(): - openai_tool: Final = { - "type": "function", - "function": { - "name": "lookup_weather", - "description": "Find current weather conditions for a location", - "parameters": { - "type": "object", - "properties": {"location": {"type": "string", "description": "City and region"}}, - "required": ["location"], - }, - }, - } - messages: Final = [{"role": "user", "content": "What's the weather?"}] - valid_tokens: Final = token_counter_new( - model="gemini-2.5-pro", - messages=messages, - tools=[openai_tool], - ) - mixed_tokens: Final = token_counter_new( - model="gemini-2.5-pro", - messages=messages, - tools=["bad", None, openai_tool], - ) - - assert mixed_tokens == valid_tokens - - -_WEATHER_FUNCTION: Final = { - "name": "get_weather", - "description": "Get the current weather for a city", - "parameters": { - "type": "object", - "properties": {"location": {"type": "string", "description": "City name"}}, - "required": ["location"], - }, -} -_TIME_FUNCTION: Final = { - "name": "get_time", - "description": "Get the local time for a city", - "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, -} -_WEB_SEARCH_FUNCTION: Final = { - "name": "web_search", - "description": "Search the web", - "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}, -} - - -def _openai_tool(function: Mapping[str, object]) -> dict[str, object]: - return {"type": "function", "function": function} - - -def _anthropic_tool(function: Mapping[str, object]) -> dict[str, object]: - return {"name": function["name"], "description": function["description"], "input_schema": function["parameters"]} - - -_USER_MESSAGES: Final = ({"role": "user", "content": "What is the weather in Paris today?"},) -_SYSTEM_AND_USER_MESSAGES: Final = ({"role": "system", "content": "You are a helpful assistant."}, *_USER_MESSAGES) - - -@pytest.mark.parametrize("tool_shape", [_openai_tool, _anthropic_tool]) -@pytest.mark.parametrize( - ("model", "messages", "functions", "billed_prompt_tokens"), - [ - # usage.prompt_tokens billed for this request by Azure AI Foundry DeepSeek-V4-Flash-0731, - # Databricks databricks-qwen35-122b-a10b and DashScope qwen3.5-plus on 2026-10-17 - ("DeepSeek-V4-Flash-0731", _USER_MESSAGES, (_WEATHER_FUNCTION,), 295), - ("DeepSeek-V4-Flash-0731", _USER_MESSAGES, (_WEATHER_FUNCTION, _TIME_FUNCTION), 348), - ("DeepSeek-V4-Flash-0731", _SYSTEM_AND_USER_MESSAGES, (_WEATHER_FUNCTION,), 300), - ("databricks-qwen35-122b-a10b", _USER_MESSAGES, (_WEATHER_FUNCTION,), 289), - ("databricks-qwen35-122b-a10b", _USER_MESSAGES, (_WEATHER_FUNCTION, _TIME_FUNCTION), 352), - ("databricks-qwen35-122b-a10b", _SYSTEM_AND_USER_MESSAGES, (_WEATHER_FUNCTION,), 296), - ("qwen3.5-plus", _USER_MESSAGES, (_WEATHER_FUNCTION,), 284), - ("qwen3.5-plus", _USER_MESSAGES, (_WEATHER_FUNCTION, _TIME_FUNCTION), 342), - ], -) -def test_token_counter_tools_match_billed_prompt_for_json_tool_chat_templates( - tool_shape: Callable[[Mapping[str, object]], dict[str, object]], - model: str, - messages: tuple[dict[str, str], ...], - functions: tuple[Mapping[str, object], ...], - billed_prompt_tokens: int, -): - counted: Final = token_counter_new( - model=model, - messages=list(messages), - tools=[tool_shape(function) for function in functions], - ) - - assert abs(counted - billed_prompt_tokens) <= billed_prompt_tokens * 0.05 - - -@pytest.mark.parametrize( - ("model", "counted_like"), - [ - ("DeepSeek-V4-Pro", "DeepSeek-V4-Flash-0731"), - ("accounts/fireworks/models/deepseek-v4-flash", "DeepSeek-V4-Flash-0731"), - ("DeepSeek-V4.1-Flash", "deepseek-v3.2"), - ("accounts/fireworks/models/deepseek-v4p1-flash", "deepseek-v3.2"), - ("Qwen/Qwen3.5-122B-A10B", "databricks-qwen35-122b-a10b"), - ("qwen3-235b-a22b", "deepseek-v3.2"), - ], -) -def test_token_counter_counts_tools_by_the_models_chat_template_family(model: str, counted_like: str): - tools: Final = [_openai_tool(_WEATHER_FUNCTION), _openai_tool(_TIME_FUNCTION)] - - assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) == token_counter_new( - model=counted_like, messages=list(_USER_MESSAGES), tools=tools - ) - - -@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"]) -@pytest.mark.parametrize( - "builtin_tool", - [ - {"type": "web_search_preview"}, - {"type": "web_search_20250305", "name": "web_search", "max_uses": 5}, - {"name": "web_search"}, - ], -) -def test_token_counter_skips_tools_not_sent_as_functions_for_chat_template_families( - model: str, builtin_tool: dict[str, object] -): - weather_tool: Final = _openai_tool(_WEATHER_FUNCTION) - - assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[builtin_tool]) == token_counter_new( - model=model, messages=list(_USER_MESSAGES) - ) - assert token_counter_new( - model=model, messages=list(_USER_MESSAGES), tools=[builtin_tool, weather_tool] - ) == token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[weather_tool]) - - -def _responses_tool(function: Mapping[str, object]) -> dict[str, object]: - return {"type": "function", **function} - - -def _gemini_tool(function: Mapping[str, object]) -> dict[str, object]: - return {"function_declarations": [function]} - - -@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"]) -@pytest.mark.parametrize("function", [_WEATHER_FUNCTION, _WEB_SEARCH_FUNCTION]) -@pytest.mark.parametrize("tool_shape", [_responses_tool, _gemini_tool]) -def test_token_counter_counts_function_tools_like_chat_completions_tools_for_chat_template_families( - model: str, function: Mapping[str, object], tool_shape: Callable[[Mapping[str, object]], dict[str, object]] -): - assert token_counter_new( - model=model, messages=list(_USER_MESSAGES), tools=[tool_shape(function)] - ) == token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[_openai_tool(function)]) - - -# Azure AI Foundry DeepSeek-V4-Flash-0731 billed 295 prompt tokens for the weather tool with each of -# these tool_choice values on 2026-10-17, the same as without one -@pytest.mark.parametrize("tool_choice", ["none", "required", {"type": "function", "function": {"name": "get_weather"}}]) -def test_token_counter_tool_choice_adds_nothing_for_deepseek_v4(tool_choice: object): - tools: Final = [_openai_tool(_WEATHER_FUNCTION)] - - assert token_counter_new( - model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools, tool_choice=tool_choice - ) == token_counter_new(model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools) - - -@pytest.mark.parametrize("model", ["DeepSeek-V4-Flash-0731", "databricks-qwen35-122b-a10b"]) -@pytest.mark.parametrize( - "anthropic_tool", - [ - {"name": "get_weather", "input_schema": _WEATHER_FUNCTION["parameters"]}, - _anthropic_tool(_WEATHER_FUNCTION), - {**_anthropic_tool(_WEATHER_FUNCTION), "strict": True}, - {**_anthropic_tool(_WEATHER_FUNCTION), "defer_loading": True, "input_examples": [{"location": "Paris"}]}, - _anthropic_tool(_WEB_SEARCH_FUNCTION), - ], -) -def test_token_counter_counts_anthropic_tools_as_sent_to_chat_template_families( - model: str, anthropic_tool: dict[str, object] -): - sent_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai([deepcopy(anthropic_tool)]) # pyright: ignore[reportArgumentType] # raw Anthropic tool dict - - assert token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[anthropic_tool]) == token_counter_new( - model=model, messages=list(_USER_MESSAGES), tools=sent_tools - ) - - -# Qwen3.5's chat_template.jinja opens a system turn for the tool schemas when the request has none, 5 tokens that a -# request with a system message already pays, while DeepSeek V4's encoding_dsv4.py appends them to the first turn -# either way (both read on Hugging Face on 2026-10-17) -@pytest.mark.parametrize( - ("model", "implicit_system_turn_tokens"), [("DeepSeek-V4-Flash-0731", 0), ("databricks-qwen35-122b-a10b", 5)] -) -def test_token_counter_counts_the_tools_system_turn_only_without_a_system_message( - model: str, implicit_system_turn_tokens: int -): - tools: Final = [_openai_tool(_WEATHER_FUNCTION)] - system_cost_with_tools: Final = token_counter_new( - model=model, messages=list(_SYSTEM_AND_USER_MESSAGES), tools=tools - ) - token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) - system_cost_without_tools: Final = token_counter_new( - model=model, messages=list(_SYSTEM_AND_USER_MESSAGES) - ) - token_counter_new(model=model, messages=list(_USER_MESSAGES)) - - assert system_cost_with_tools == system_cost_without_tools - implicit_system_turn_tokens - - class NeedsToleranceUpdateError(Exception): """Custom exception to mark tests that have improved""" diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index fea65a9c821..ddac561f337 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -1,113 +1,7 @@ -import asyncio -from base64 import b64encode -from copy import deepcopy -from threading import get_ident -from typing import Final -import httpx -import pytest -import respx -from pydantic import JsonValue, TypeAdapter - -import litellm -from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, ) -from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( - AzureAIAnthropicCountTokensHandler, -) -from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( - AzureAIAnthropicCountTokensConfig, -) - - -@pytest.mark.parametrize( - "config_type", (AnthropicCountTokensConfig, AzureAIAnthropicCountTokensConfig) -) -@pytest.mark.parametrize( - ("image_url", "source"), - ( - ("data:image/png;base64,aW1hZ2U=", {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}), - ({"url": "data:image/png;base64,aW1hZ2U="}, {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}), - ( - {"url": "data:image/png;base64,aW1hZ2U=", "format": "image/jpeg", "detail": "high"}, - {"type": "base64", "media_type": "image/jpeg", "data": "aW1hZ2U="}, - ), - ), -) -def test_count_translates_openai_images_without_mutating_input( - config_type: type[AnthropicCountTokensConfig], - image_url: str | dict[str, JsonValue], - source: dict[str, JsonValue], -) -> None: - cache_control: Final[dict[str, JsonValue]] = {"type": "ephemeral"} - messages: Final[list[dict[str, JsonValue]]] = [{ - "role": "user", "content": [ - {"type": "text", "text": "Count this image"}, - {"type": "image_url", "image_url": image_url, "cache_control": cache_control}, - ], - }] - original: Final = deepcopy(messages) - result: Final = config_type().transform_request_to_count_tokens( - model="claude-opus-5-5", messages=messages - ) - - assert result == { - "model": "claude-opus-5-5", "messages": [{ - "role": "user", "content": [ - {"type": "text", "text": "Count this image"}, - {"type": "image", "source": source, "cache_control": cache_control}, - ], - }], - } - assert messages == original - - -@pytest.mark.parametrize( - "config_type", (AnthropicCountTokensConfig, AzureAIAnthropicCountTokensConfig) -) -def test_count_normalizes_nested_tool_images_and_preserves_native_fields( - config_type: type[AnthropicCountTokensConfig], -) -> None: - openai_image: Final[dict[str, JsonValue]] = { - "type": "image_url", "image_url": {"url": "data:image/png;base64,aW1hZ2U="} - } - native_image: Final[dict[str, JsonValue]] = { - "type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="} - } - assistant: Final[dict[str, JsonValue]] = {"role": "assistant", "content": [ - {"type": "thinking", "thinking": "inspect screenshot", "signature": "fixture-signature"}, - {"type": "tool_use", "id": "read-1", "name": "Read", "input": {"content": [openai_image]}}, - ]} - tool_result: Final[dict[str, JsonValue]] = { - "type": "tool_result", "tool_use_id": "read-1", "is_error": False, - "content": [{"type": "text", "text": "Screenshot"}, native_image, openai_image], - "cache_control": {"type": "ephemeral"}, - } - text_result: Final[dict[str, JsonValue]] = {"type": "tool_result", "tool_use_id": "read-2", "content": "done"} - messages: Final[list[dict[str, JsonValue]]] = [ - assistant, {"role": "user", "content": [native_image, tool_result, text_result]} - ] - tools: Final[list[dict[str, JsonValue]]] = [{ - "name": "Read", "input_schema": {"type": "object", "examples": [openai_image]} - }] - system: Final[JsonValue] = [{"type": "text", "text": "policy", "cache_control": {"type": "ephemeral"}}] - options: Final[dict[str, JsonValue]] = { - "thinking": {"type": "adaptive"}, "tool_choice": {"type": "auto"}, "output_config": {"effort": "high"} - } - original: Final = deepcopy((messages, tools, system, options)) - result: Final = config_type().transform_request_to_count_tokens( - model="claude-opus-5-5", messages=messages, tools=tools, system=system, optional_params=options - ) - - assert result == { - "model": "claude-opus-5-5", "system": system, "tools": tools, **options, - "messages": [assistant, {"role": "user", "content": [native_image, { - **tool_result, "content": [{"type": "text", "text": "Screenshot"}, native_image, native_image] - }, text_result]}], - } - assert (messages, tools, system, options) == original def test_transform_basic_request(): @@ -194,188 +88,3 @@ def test_transform_no_system_no_tools(): assert "system" not in result assert "tools" not in result - - -@pytest.fixture -def httpx_transport_clients(monkeypatch): - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - client_cache = getattr(litellm, "in_memory_llm_clients_cache", None) - if client_cache is not None: - client_cache.flush_cache() - yield - if client_cache is not None: - client_cache.flush_cache() - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("httpx_transport_clients") -@pytest.mark.parametrize( - "handler_type", (AnthropicCountTokensHandler, AzureAIAnthropicCountTokensHandler) -) -@pytest.mark.parametrize("scheme", ("http", "https")) -@pytest.mark.parametrize("dict_url", (False, True)) -async def test_remote_image_fetch_keeps_counting_handler_event_loop_responsive( - handler_type: type[AnthropicCountTokensHandler] | type[AzureAIAnthropicCountTokensHandler], - scheme: str, - dict_url: bool, -) -> None: - loop: Final = asyncio.get_running_loop() - loop_thread: Final = get_ident() - witness: Final = asyncio.Event() - image_bytes: Final = b"\x89PNG\r\n\x1a\ncount-image" - image_url: Final = f"{scheme}://1.1.1.1/{handler_type.__name__}-{dict_url}.png" - model: Final = "claude-opus-5-5" - api_base: Final = "https://gateway.example/anthropic" - image: Final[dict[str, JsonValue]] = { - "type": "image_url", "image_url": {"url": image_url} if dict_url else image_url, - "cache_control": {"type": "ephemeral"}, - } - messages: Final[list[dict[str, JsonValue]]] = [{ - "role": "user", "content": [image, {"type": "tool_result", "tool_use_id": "read-1", "content": [image]}] - }] - original: Final = deepcopy(messages) - - async def run_witness() -> None: - witness.set() - - def image_response(_request: httpx.Request) -> httpx.Response: - assert get_ident() != loop_thread, "image fetch blocked the counting handler's event loop" - asyncio.run_coroutine_threadsafe(run_witness(), loop).result(timeout=5) - return httpx.Response(200, content=image_bytes, headers={"Content-Type": "image/png"}) - - with respx.mock: - image_route: Final = respx.get(image_url).mock(side_effect=image_response) - count_route: Final = respx.post(url__startswith=api_base).mock( - return_value=httpx.Response(200, json={"input_tokens": 7}) - ) - handler: Final = handler_type() - result: Final = await handler.handle_count_tokens_request( - model=model, messages=messages, api_base=api_base, api_key="test-key" - ) - - assert witness.is_set() - assert image_route.call_count == count_route.call_count == 1 - assert result == {"input_tokens": 7} - native_image: Final[dict[str, JsonValue]] = { - "type": "image", "source": { - "type": "base64", "media_type": "image/png", "data": b64encode(image_bytes).decode() - }, "cache_control": {"type": "ephemeral"}, - } - assert TypeAdapter(dict[str, JsonValue]).validate_json(count_route.calls.last.request.content) == { - "model": model, "messages": [{"role": "user", "content": [ - native_image, {"type": "tool_result", "tool_use_id": "read-1", "content": [native_image]} - ]}], - } - assert messages == original - - -@pytest.mark.parametrize( - "config_type", (AnthropicCountTokensConfig, AzureAIAnthropicCountTokensConfig) -) -def test_count_lifts_the_leading_system_run_into_system( - config_type: type[AnthropicCountTokensConfig], -) -> None: - """A Responses ``instructions`` arrives as a leading system-role message. count_tokens answers 400 - on that role at the head of ``messages`` and only takes the initial prompt in ``system``, so the - leading run moves there with its cache_control, empty text dropped, and a later reminder stays.""" - cache_control: Final[dict[str, JsonValue]] = {"type": "ephemeral"} - messages: Final[list[dict[str, JsonValue]]] = [ - {"role": "system", "content": "Be terse", "cache_control": cache_control}, - {"role": "system", "content": [{"type": "text", "text": "Answer in French"}, {"type": "text", "text": ""}]}, - {"role": "user", "content": "Hello, how are you?"}, - {"role": "system", "content": "later reminder"}, - {"role": "assistant", "content": "Bonjour."}, - ] - original: Final = deepcopy(messages) - result: Final = config_type().transform_request_to_count_tokens(model="claude-opus-5-5", messages=messages) - - assert result == { - "model": "claude-opus-5-5", - "system": [ - {"type": "text", "text": "Be terse", "cache_control": cache_control}, - {"type": "text", "text": "Answer in French"}, - ], - "messages": [ - {"role": "user", "content": "Hello, how are you?"}, - {"role": "system", "content": "later reminder"}, - {"role": "assistant", "content": "Bonjour."}, - ], - } - assert messages == original - - -@pytest.mark.parametrize( - ("system", "expected_system"), - ( - (None, [{"type": "text", "text": "Be terse"}]), - ("", [{"type": "text", "text": "Be terse"}]), - ("Answer in French", [{"type": "text", "text": "Answer in French"}, {"type": "text", "text": "Be terse"}]), - ( - [{"type": "text", "text": "Answer in French", "cache_control": {"type": "ephemeral"}}], - [ - {"type": "text", "text": "Answer in French", "cache_control": {"type": "ephemeral"}}, - {"type": "text", "text": "Be terse"}, - ], - ), - ), - ids=["absent", "empty", "string", "blocks"], -) -def test_count_keeps_the_callers_system_ahead_of_the_lifted_run( - system: JsonValue, expected_system: list[dict[str, JsonValue]] -) -> None: - result: Final = AnthropicCountTokensConfig().transform_request_to_count_tokens( - model="claude-opus-5-5", - messages=[{"role": "system", "content": "Be terse"}, {"role": "user", "content": "hi"}], - system=system, - ) - - assert result == { - "model": "claude-opus-5-5", - "system": expected_system, - "messages": [{"role": "user", "content": "hi"}], - } - - -def test_count_leaves_a_non_text_system_and_its_messages_as_sent() -> None: - """A malformed ``system`` is the provider's to reject, so nothing is rearranged around it.""" - messages: Final[list[dict[str, JsonValue]]] = [ - {"role": "system", "content": "Be terse"}, - {"role": "user", "content": "hi"}, - ] - result: Final = AnthropicCountTokensConfig().transform_request_to_count_tokens( - model="claude-opus-5-5", messages=messages, system=5 - ) - - assert result == {"model": "claude-opus-5-5", "system": 5, "messages": messages} - - -def test_count_drops_a_leading_system_message_without_text() -> None: - result: Final = AnthropicCountTokensConfig().transform_request_to_count_tokens( - model="claude-opus-5-5", - messages=[{"role": "system", "content": ""}, {"role": "user", "content": "hi"}], - ) - - assert result == {"model": "claude-opus-5-5", "messages": [{"role": "user", "content": "hi"}]} - - -@pytest.mark.asyncio -async def test_handler_sends_the_leading_system_run_as_system_not_as_a_message(httpx_transport_clients): - """The wire body is what the provider judges: ``system`` carries the prompt and no message has - ``role: "system"``, so a Responses ``instructions`` is counted by Anthropic instead of 400ing.""" - with respx.mock: - route = respx.post(url__startswith="https://gateway.example").mock( - return_value=httpx.Response(200, json={"input_tokens": 21}) - ) - result = await AnthropicCountTokensHandler().handle_count_tokens_request( - model="claude-opus-5-5", - messages=[{"role": "system", "content": "Be terse"}, {"role": "user", "content": "Hello, how are you?"}], - api_key="sk-ant-api03-test-key", - api_base="https://gateway.example", - ) - - assert result == {"input_tokens": 21} - assert TypeAdapter(dict[str, JsonValue]).validate_json(route.calls.last.request.content) == { - "model": "claude-opus-5-5", - "system": [{"type": "text", "text": "Be terse"}], - "messages": [{"role": "user", "content": "Hello, how are you?"}], - } diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py deleted file mode 100644 index 71b451f834c..00000000000 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py +++ /dev/null @@ -1,109 +0,0 @@ -import json -from collections.abc import Mapping -from typing import Final - -import httpx -import pytest - -from litellm.llms.bedrock.common_utils import BedrockError -from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - -MANTLE_COUNT_URL: Final = "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages/count_tokens" -LITELLM_PARAMS: Final = { - "aws_access_key_id": "AKIATESTACCESSKEY", - "aws_secret_access_key": "test-secret", - "aws_region_name": "us-east-1", -} -REQUEST: Final[dict[str, object]] = { - "model": "global.anthropic.claude-opus-4-8", - "messages": [{"role": "user", "content": "The quick brown fox"}], - "system": "You are a terse assistant.", - "tools": [{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}], -} - - -class _MantleEndpoint: - def __init__(self, status_code: int, body: Mapping[str, object]) -> None: - self.status_code: Final = status_code - self.body: Final = body - self.requests: tuple[httpx.Request, ...] = () - - def __call__(self, request: httpx.Request) -> httpx.Response: - self.requests = (*self.requests, request) - return httpx.Response(self.status_code, json=dict(self.body), request=request) - - def only_request(self) -> httpx.Request: - assert len(self.requests) == 1, self.requests - return self.requests[0] - - -def _client(endpoint: _MantleEndpoint) -> AsyncHTTPHandler: - return AsyncHTTPHandler(transport=httpx.MockTransport(endpoint)) - - -@pytest.fixture(autouse=True) -def _sigv4_only(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) - monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) - - -@pytest.mark.asyncio -async def test_counts_on_mantle_with_the_base_model_and_the_deployment_credentials() -> None: - mantle: Final = _MantleEndpoint(200, {"input_tokens": 2177}) - - result: Final = await BedrockMantleCountTokensHandler().handle_count_tokens_request( - request_data=dict(REQUEST), - litellm_params=dict(LITELLM_PARAMS), - resolved_model="anthropic.claude-opus-4-8", - client=_client(mantle), - ) - - assert result == {"input_tokens": 2177} - posted: Final = mantle.only_request() - assert str(posted.url) == MANTLE_COUNT_URL - assert posted.headers["Authorization"].startswith("AWS4-HMAC-SHA256") - assert "us-east-1/bedrock/aws4_request" in posted.headers["Authorization"] - assert posted.headers["anthropic-version"] == "2023-06-01" - assert json.loads(posted.content) == { - "model": "anthropic.claude-opus-4-8", - "messages": REQUEST["messages"], - "system": REQUEST["system"], - "tools": REQUEST["tools"], - } - - -@pytest.mark.asyncio -async def test_bedrock_mantle_api_base_env_names_the_host(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", "https://vpce-abc.bedrock-mantle.us-east-1.vpce.api.aws") - mantle: Final = _MantleEndpoint(200, {"input_tokens": 3}) - - await BedrockMantleCountTokensHandler().handle_count_tokens_request( - request_data=dict(REQUEST), - litellm_params=dict(LITELLM_PARAMS), - resolved_model="anthropic.claude-opus-4-8", - client=_client(mantle), - ) - - assert ( - str(mantle.only_request().url) - == "https://vpce-abc.bedrock-mantle.us-east-1.vpce.api.aws/anthropic/v1/messages/count_tokens" - ) - - -@pytest.mark.asyncio -async def test_non_200_answers_raise_bedrock_error_with_mantle_status() -> None: - mantle: Final = _MantleEndpoint( - 404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}} - ) - - with pytest.raises(BedrockError) as raised: - await BedrockMantleCountTokensHandler().handle_count_tokens_request( - request_data=dict(REQUEST), - litellm_params=dict(LITELLM_PARAMS), - resolved_model="anthropic.claude-sonnet-5-5", - client=_client(mantle), - ) - - assert raised.value.status_code == 404 - assert "does not exist" in raised.value.message diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py deleted file mode 100644 index 8a21027872a..00000000000 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ /dev/null @@ -1,176 +0,0 @@ -from collections.abc import Mapping -from typing import Final - -import httpx -import pytest -from pydantic import JsonValue, TypeAdapter - -from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - -RUNTIME_HOST: Final = "bedrock-runtime.us-east-1.amazonaws.com" -MANTLE_HOST: Final = "bedrock-mantle.us-east-1.api.aws" -UNSUPPORTED: Final = {"message": "The provided model doesn't support counting tokens."} -DEPLOYMENT: Final = { - "litellm_params": { - "aws_access_key_id": "AKIATESTACCESSKEY", - "aws_secret_access_key": "test-secret", - "aws_region_name": "us-east-1", - } -} -MESSAGES: Final = [{"role": "user", "content": "The quick brown fox jumps over the lazy dog."}] -_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue]) - - -class _Bedrock: - def __init__(self, runtime: tuple[int, Mapping[str, object]], mantle: tuple[int, Mapping[str, object]]) -> None: - self.runtime: Final = runtime - self.mantle: Final = mantle - self.requests: tuple[httpx.Request, ...] = () - - def __call__(self, request: httpx.Request) -> httpx.Response: - self.requests = (*self.requests, request) - status_code, body = self.runtime if request.url.host == RUNTIME_HOST else self.mantle - return httpx.Response(status_code, json=dict(body), request=request) - - def posted_hosts(self) -> tuple[str, ...]: - return tuple(request.url.host for request in self.requests) - - -def _counter(bedrock: _Bedrock) -> BedrockTokenCounter: - return BedrockTokenCounter(client=AsyncHTTPHandler(transport=httpx.MockTransport(bedrock))) - - -@pytest.fixture(autouse=True) -def _sigv4_only(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) - monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) - - -@pytest.mark.asyncio -async def test_claude_model_bedrock_runtime_cannot_count_is_counted_on_mantle() -> None: - bedrock: Final = _Bedrock(runtime=(400, UNSUPPORTED), mantle=(200, {"input_tokens": 2177})) - - result: Final = await _counter(bedrock).count_tokens( - model_to_use="global.anthropic.claude-opus-4-8", - messages=MESSAGES, - contents=None, - deployment=DEPLOYMENT, - request_model="claude-opus-4-8", - system="You are a terse assistant.", - ) - - assert result is not None - assert result.error is False - assert result.total_tokens == 2177 - assert result.tokenizer_type == "bedrock_mantle_api" - assert result.original_response == {"input_tokens": 2177} - assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) - mantle_body: Final = _JSON_BODY.validate_json(bedrock.requests[1].content) - assert mantle_body["model"] == "anthropic.claude-opus-4-8" - assert mantle_body["system"] == "You are a terse assistant." - - -@pytest.mark.asyncio -async def test_bedrock_runtime_count_is_kept_when_it_answers() -> None: - bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336})) - - result: Final = await _counter(bedrock).count_tokens( - model_to_use="global.anthropic.claude-sonnet-4-6", - messages=MESSAGES, - contents=None, - deployment=DEPLOYMENT, - request_model="claude-sonnet-4-6", - ) - - assert result is not None - assert result.total_tokens == 1353 - assert result.tokenizer_type == "bedrock_api" - assert bedrock.posted_hosts() == (RUNTIME_HOST,) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("model_to_use", "runtime"), - ( - ("global.anthropic.claude-opus-4-8", (401, {"Message": "not authorized to perform: bedrock:CountTokens"})), - ("global.anthropic.claude-opus-4-8", (403, {"Message": "not authorized to perform: bedrock:CountTokens"})), - ("amazon.nova-pro-v1:0", (400, UNSUPPORTED)), - ("amazon.nova-pro-v1:0", (404, UNSUPPORTED)), - ), -) -async def test_other_bedrock_runtime_failures_are_not_retried_on_mantle( - model_to_use: str, runtime: tuple[int, Mapping[str, object]] -) -> None: - bedrock: Final = _Bedrock(runtime=runtime, mantle=(200, {"input_tokens": 2177})) - - result: Final = await _counter(bedrock).count_tokens( - model_to_use=model_to_use, - messages=MESSAGES, - contents=None, - deployment=DEPLOYMENT, - request_model=model_to_use, - ) - - assert result is not None - assert result.error is True - assert result.status_code == runtime[0] - assert result.tokenizer_type == "bedrock_api" - assert bedrock.posted_hosts() == (RUNTIME_HOST,) - - -@pytest.mark.asyncio -async def test_mantle_failure_is_reported_with_its_status() -> None: - bedrock: Final = _Bedrock( - runtime=(404, UNSUPPORTED), - mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}), - ) - - result: Final = await _counter(bedrock).count_tokens( - model_to_use="global.anthropic.claude-sonnet-5-5", - messages=MESSAGES, - contents=None, - deployment=DEPLOYMENT, - request_model="claude-sonnet-5-5", - ) - - assert result is not None - assert result.error is True - assert result.status_code == 404 - assert result.tokenizer_type == "bedrock_mantle_api" - assert "does not exist" in (result.error_message or "") - assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) - - -@pytest.mark.asyncio -async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> None: - bedrock: Final = _Bedrock(runtime=(404, UNSUPPORTED), mantle=(200, {"input_tokens": 2177})) - tools: Final = [ - { - "name": "lookup", - "description": "Find a record", - "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, - } - ] - - result: Final = await _counter(bedrock).count_tokens( - model_to_use="global.anthropic.claude-opus-4-8", - messages=MESSAGES, - contents=None, - deployment=DEPLOYMENT, - request_model="claude-opus-4-8", - system="You are a terse assistant.", - tools=tools, - ) - - assert result is not None - assert result.error is False - assert result.total_tokens == 2177 - assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) - mantle_body: Final = _JSON_BODY.validate_json(bedrock.requests[1].content) - assert mantle_body == { - "model": "anthropic.claude-opus-4-8", - "messages": MESSAGES, - "system": "You are a terse assistant.", - "tools": tools, - } diff --git a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py index ddcbfc61822..77b8e66c5a8 100644 --- a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -83,54 +83,6 @@ def test_transform_includes_tools(): assert result["tools"] == tools -_SCHEMA: Final[dict[str, object]] = {"type": "object", "properties": {}} # mutable-ok: matches the tools param type - - -@pytest.mark.parametrize( - ("tool", "expected"), - ( - ( - {"name": "lookup", "input_schema": _SCHEMA}, - {"type": "function", "name": "lookup", "description": "", "parameters": _SCHEMA, "strict": False}, - ), - ( - {"name": "lookup", "input_schema": _SCHEMA, "strict": True}, - {"type": "function", "name": "lookup", "description": "", "parameters": _SCHEMA, "strict": True}, - ), - ( - {"type": "custom", "name": "lookup", "description": "Find a record", "input_schema": _SCHEMA}, - { - "type": "function", - "name": "lookup", - "description": "Find a record", - "parameters": _SCHEMA, - "strict": False, - }, - ), - ), - ids=("anthropic", "anthropic-strict", "anthropic-custom-type"), -) -def test_transform_anthropic_tool_shapes(tool: dict[str, object], expected: dict[str, object]) -> None: - result: Final = OpenAICountTokensConfig().transform_request_to_count_tokens( - model="gpt-4o", input="Find a record", tools=[tool] - ) - - assert result["tools"] == [expected] - - -def test_transform_builtin_tool_passes_through(): - config: Final = OpenAICountTokensConfig() - tools: Final = [{"type": "web_search", "search_context_size": "low"}] - - result: Final = config.transform_request_to_count_tokens( - model="gpt-4o", - input="Find a record", - tools=tools, - ) - - assert result["tools"] == tools - - def test_transform_no_instructions_no_tools(): """Test that None values are not included.""" config = OpenAICountTokensConfig() @@ -154,7 +106,9 @@ def test_messages_to_responses_input_basic(): {"role": "user", "content": "How are you?"}, ] - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( + messages + ) assert len(input_items) == 3 assert input_items[0] == {"role": "user", "content": "Hello"} @@ -170,7 +124,9 @@ def test_messages_to_responses_input_with_system(): {"role": "user", "content": "Hello"}, ] - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( + messages + ) assert len(input_items) == 1 assert input_items[0] == {"role": "user", "content": "Hello"} @@ -184,7 +140,9 @@ def test_messages_to_responses_input_with_developer(): {"role": "user", "content": "Hello"}, ] - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( + messages + ) assert len(input_items) == 1 assert instructions == "Be concise." @@ -197,7 +155,9 @@ def test_messages_to_responses_input_with_tool(): {"role": "tool", "content": "72°F", "tool_call_id": "call_123"}, ] - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( + messages + ) assert len(input_items) == 2 assert input_items[1] == { @@ -468,7 +428,10 @@ def test_validate_request_missing_input(): def test_get_endpoint_default(): """Test default endpoint URL.""" config = OpenAICountTokensConfig() - assert config.get_openai_count_tokens_endpoint() == "https://api.openai.com/v1/responses/input_tokens" + assert ( + config.get_openai_count_tokens_endpoint() + == "https://api.openai.com/v1/responses/input_tokens" + ) def test_get_endpoint_custom_base(): diff --git a/tests/unit/llms/vertex_ai/count_tokens/__init__.py b/tests/unit/llms/vertex_ai/count_tokens/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 71e6ecc1f81..f9561a0d710 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -47,35 +47,6 @@ async def test_vertex_gemini_rejects_non_text_system_blocks_before_request( assert "text blocks" in (result.error_message or "") -@pytest.mark.asyncio -@pytest.mark.parametrize( - "model", - ( - "zai-org/glm-5.2-maas", - "deepseek-ai/deepseek-v3.2-maas", - "qwen/qwen3-coder-480b-a35b-instruct-maas", - "mistral-medium-3", - "codestral-2501", - "openai/gpt-oss-120b-maas", - "google/gemma-4-26b-a4b-it-maas", - ), -) -async def test_vertex_partner_model_without_an_inference_matching_count_answers_400_naming_the_model(model: str) -> None: - result: Final = await VertexAITokenCounter().count_tokens( - model_to_use=model, - messages=[{"role": "user", "content": "Count this"}], - contents=None, - deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "global"}}, - request_model=f"vertex_ai/{model}", - messages_format="anthropic", - ) - - assert result is not None - assert result.error is True - assert result.status_code == 400 - assert result.error_message == f"LiteLLM does not support counting tokens for Vertex AI model {model}." - - class _VertexCountTokensPeer: def __init__(self, status_code: int) -> None: self.status_code: Final = status_code diff --git a/tests/unit/proxy/test_custom_tokenizer_bug.py b/tests/unit/proxy/test_custom_tokenizer_bug.py index 29625ea1a32..23b6b302202 100644 --- a/tests/unit/proxy/test_custom_tokenizer_bug.py +++ b/tests/unit/proxy/test_custom_tokenizer_bug.py @@ -9,7 +9,6 @@ The HuggingFace download boundary (Tokenizer.from_pretrained) is mocked so these stay hermetic unit tests; the proxy's extraction-and-selection path runs for real. """ -from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -18,7 +17,7 @@ import litellm import litellm.proxy.proxy_server import litellm.utils from litellm import Router -from litellm.proxy._types import ProxyException, TokenCountRequest +from litellm.proxy._types import TokenCountRequest from litellm.proxy.proxy_server import token_counter @@ -103,60 +102,3 @@ async def test_model_without_custom_tokenizer_uses_default(monkeypatch): assert response.tokenizer_type == "openai_tokenizer" assert response.model_used == "gpt-4" assert response.total_tokens > 0 - - - -def _vertex_glm_router_with_its_own_tokenizer() -> Router: - return Router( - model_list=[ - { - "model_name": "vertex-glm", - "litellm_params": { - "model": "vertex_ai/zai-org/glm-5.2-maas", - "vertex_project": "test-project", - "vertex_location": "global", - }, - "model_info": { - "custom_tokenizer": {"identifier": "zai-org/GLM-5.2", "revision": "main", "auth_token": None}, - }, - } - ] - ) - - -@pytest.mark.asyncio -async def test_custom_tokenizer_counts_a_vertex_partner_model_without_a_count_route( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _vertex_glm_router_with_its_own_tokenizer()) - - with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls: - mock_tokenizer_cls.from_pretrained.return_value = _fake_hf_tokenizer(9) - - response: Final = await token_counter( - request=TokenCountRequest(model="vertex-glm", messages=[{"role": "user", "content": "Return exactly OK."}]), - call_endpoint=True, - ) - - mock_tokenizer_cls.from_pretrained.assert_called_once_with("zai-org/GLM-5.2", revision="main", token=None) - assert response.error is not True - assert response.tokenizer_type == "huggingface_tokenizer" - assert response.total_tokens >= 9 - - -@pytest.mark.asyncio -async def test_contents_only_count_of_a_vertex_partner_model_without_a_count_route_answers_400( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _vertex_glm_router_with_its_own_tokenizer()) - - with pytest.raises(ProxyException) as raised: - await token_counter( - request=TokenCountRequest( - model="vertex-glm", contents=[{"role": "user", "parts": [{"text": "Return exactly OK."}]}] - ), - call_endpoint=True, - ) - - assert raised.value.code == "400" - assert raised.value.message == "LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas." diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 86556f1c7bd..38f2181d4e7 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -801,45 +801,41 @@ def test_vertex_ai_partner_models_token_counting_endpoint(vertex_location): assert endpoint.startswith(f"https://{vertex_location}-aiplatform.googleapis.com") -class _FailingRuntimeHandler(BedrockCountTokensHandler): - def __init__(self, error: Exception) -> None: - super().__init__() - self._error = error - - async def handle_count_tokens_request( - self, - request_data: dict[str, object], - litellm_params: dict[str, object], - resolved_model: str, - client: AsyncHTTPHandler | None = None, - ) -> dict[str, JsonValue]: - raise self._error - - @pytest.mark.asyncio async def test_bedrock_token_counter_error_propagation_bedrock_error(): """ Test that BedrockTokenCounter properly returns error response when BedrockError is raised. Verifies that the status code and error message are preserved. """ - counter = BedrockTokenCounter( - runtime_handler=_FailingRuntimeHandler(BedrockError(status_code=429, message="Rate limit exceeded")) - ) + counter = BedrockTokenCounter() - result = await counter.count_tokens( - model_to_use="anthropic.claude-3-sonnet", - messages=[{"role": "user", "content": "hello"}], - contents=None, - deployment={"litellm_params": {}}, - request_model="bedrock/anthropic.claude-3-sonnet", - ) + # Mock the handler to raise BedrockError with specific status code + with patch.object( + counter, "count_tokens", wraps=counter.count_tokens + ) as mock_count: + # We need to patch at the handler level + with patch( + "litellm.llms.bedrock.count_tokens.bedrock_token_counter.BedrockCountTokensHandler" + ) as MockHandler: + mock_handler_instance = MockHandler.return_value + mock_handler_instance.handle_count_tokens_request = AsyncMock( + side_effect=BedrockError(status_code=429, message="Rate limit exceeded") + ) - assert result is not None - assert result.error is True - assert result.status_code == 429 - assert "Rate limit exceeded" in result.error_message - assert result.tokenizer_type == "bedrock_api" - assert result.total_tokens == 0 + result = await counter.count_tokens( + model_to_use="anthropic.claude-3-sonnet", + messages=[{"role": "user", "content": "hello"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="bedrock/anthropic.claude-3-sonnet", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 429 + assert "Rate limit exceeded" in result.error_message + assert result.tokenizer_type == "bedrock_api" + assert result.total_tokens == 0 @pytest.mark.asyncio @@ -847,20 +843,28 @@ async def test_bedrock_token_counter_error_propagation_generic_exception(): """ Test that BedrockTokenCounter returns error response with 500 status for generic exceptions. """ - counter = BedrockTokenCounter(runtime_handler=_FailingRuntimeHandler(Exception("Unexpected error"))) + counter = BedrockTokenCounter() - result = await counter.count_tokens( - model_to_use="anthropic.claude-3-sonnet", - messages=[{"role": "user", "content": "hello"}], - contents=None, - deployment={"litellm_params": {}}, - request_model="bedrock/anthropic.claude-3-sonnet", - ) + with patch( + "litellm.llms.bedrock.count_tokens.bedrock_token_counter.BedrockCountTokensHandler" + ) as MockHandler: + mock_handler_instance = MockHandler.return_value + mock_handler_instance.handle_count_tokens_request = AsyncMock( + side_effect=Exception("Unexpected error") + ) - assert result is not None - assert result.error is True - assert result.status_code == 500 - assert "Unexpected error" in result.error_message + result = await counter.count_tokens( + model_to_use="anthropic.claude-3-sonnet", + messages=[{"role": "user", "content": "hello"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="bedrock/anthropic.claude-3-sonnet", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 500 + assert "Unexpected error" in result.error_message @pytest.mark.asyncio diff --git a/tests/unit/test_cost_map_guard.py b/tests/unit/test_cost_map_guard.py index 410b3acfd42..e2c54d381d8 100644 --- a/tests/unit/test_cost_map_guard.py +++ b/tests/unit/test_cost_map_guard.py @@ -22,6 +22,7 @@ def _load(name: str) -> ModuleType: return module +schema_module: Final = _load("generate_model_prices_schema") guard: Final = _load("cost_map_guard") MAP_FILES: Final = (guard.COST_MAP_PATH,) @@ -51,9 +52,12 @@ def _serialize(cost_map: dict[str, object]) -> str: return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n" -def _snapshot(cost_map: dict[str, object], backup: str | None = None) -> object: +def _snapshot(cost_map: dict[str, object], backup: str | None = None, schema: str | None = None) -> object: text = _serialize(cost_map) - return guard.Snapshot(cost_map=text, backup=text if backup is None else backup) + rendered = schema_module.render(schema_module.build_schema(cost_map)) + return guard.Snapshot( + cost_map=text, backup=text if backup is None else backup, schema=rendered if schema is None else schema + ) BASE: Final = _snapshot(BASE_MAP) @@ -131,14 +135,14 @@ def test_bot_may_add_and_reprice_models() -> None: def test_broken_json_is_reported() -> None: - head = guard.Snapshot(cost_map="{not json", backup="{not json") + head = guard.Snapshot(cost_map="{not json", backup="{not json", schema="{}") assert _failures(head, bot=False) == ( f"{guard.COST_MAP_PATH} is not valid JSON: Expecting property name enclosed in double quotes: line 1 column 2 (char 1)", ) def test_non_object_root_is_reported() -> None: - head = guard.Snapshot(cost_map="[]", backup="[]") + head = guard.Snapshot(cost_map="[]", backup="[]", schema="{}") assert _failures(head, bot=False) == (f"{guard.COST_MAP_PATH} must be a JSON object at the root",) @@ -147,7 +151,28 @@ def test_backup_drift_is_reported() -> None: assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.BACKUP_PATH)] -STALE_HEAD: Final = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)})) +def test_schema_out_of_sync_is_reported() -> None: + head = _snapshot({**BASE_MAP, "openrouter/c": _entry(supports_audio_input=True)}, schema=BASE.schema) + assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.SCHEMA_PATH)] + + +def test_schema_validation_errors_are_reported() -> None: + head = _snapshot({**BASE_MAP, "openrouter/c": _entry(-1e-06)}) + prefix = f"{guard.COST_MAP_PATH} does not validate against its schema: openrouter/c." + assert [failure.removeprefix(prefix).split(":")[0] for failure in _failures(head, bot=False)] == [ + "input_cost_per_token", + "output_cost_per_token", + ] + + +def test_unclassified_entry_key_is_reported() -> None: + text = _serialize({**BASE_MAP, "openrouter/c": _entry(weird_thing=1)}) + head = guard.Snapshot(cost_map=text, backup=text, schema=BASE.schema) + (failure,) = _failures(head, bot=False) + assert "Unclassified keys" in failure and "weird_thing" in failure + + +STALE_HEAD: Final = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)}), schema="{}") CODE_ONLY: Final = ( "litellm/utils.py", "tests/test_litellm/test_utils.py", @@ -156,7 +181,7 @@ CODE_ONLY: Final = ( def test_human_pr_that_leaves_the_cost_map_alone_skips_the_file_checks() -> None: - unparseable: Final = guard.Snapshot(cost_map="{not json", backup="") + unparseable: Final = guard.Snapshot(cost_map="{not json", backup="", schema="") assert _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=False) == () assert _failures(STALE_HEAD, changed_files=(), bot=False) == () assert _failures(unparseable, changed_files=CODE_ONLY, bot=False) == () @@ -166,16 +191,18 @@ def test_human_pr_that_leaves_the_cost_map_alone_skips_the_file_checks() -> None def test_touching_any_cost_map_file_keeps_the_file_checks(guarded_path: str) -> None: failures: Final = _failures(STALE_HEAD, changed_files=(*CODE_ONLY, guarded_path), bot=False) assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)] + assert [failure for failure in failures if failure.startswith(guard.SCHEMA_PATH)] def test_bot_pr_always_gets_the_file_checks() -> None: failures: Final = _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=True) assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)] + assert [failure for failure in failures if failure.startswith(guard.SCHEMA_PATH)] def test_contract_names_the_skip() -> None: assert guard.contract_for(False, CODE_ONLY) == "human PR, cost map untouched" - assert guard.contract_for(False, (*CODE_ONLY, guard.COST_MAP_PATH)) == "human PR, file checks only" + assert guard.contract_for(False, (*CODE_ONLY, guard.SCHEMA_PATH)) == "human PR, file checks only" assert guard.contract_for(True, CODE_ONLY) == "bot contract enforced" @@ -188,13 +215,6 @@ def test_bot_may_only_touch_the_cost_map_files() -> None: ) -def test_bot_may_update_cost_map_and_legacy_schema() -> None: - head: Final = _snapshot({**BASE_MAP, "openrouter/c": _entry()}) - changed: Final = (*guard.GUARDED_PATHS, guard.LEGACY_SCHEMA_PATH) - - assert _failures(head, changed_files=changed) == () - - def test_bot_may_not_remove_models() -> None: head = _snapshot({key: value for key, value in BASE_MAP.items() if key != "openrouter/b"}) assert _failures(head, bot=False) == () @@ -218,6 +238,7 @@ def _commit(repo: Path, cost_map: dict[str, object], message: str) -> str: (repo / guard.COST_MAP_PATH).write_text(text) (repo / guard.BACKUP_PATH).parent.mkdir(exist_ok=True) (repo / guard.BACKUP_PATH).write_text(text) + (repo / guard.SCHEMA_PATH).write_text(schema_module.render(schema_module.build_schema(cost_map))) return _git_commit(repo, message) @@ -271,7 +292,8 @@ def test_main_skips_the_file_checks_on_a_stale_base_the_pr_never_touched(tmp_pat subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) _commit(tmp_path, BASE_MAP, "base") (tmp_path / guard.BACKUP_PATH).write_text(_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)})) - stale_base: Final = _commit_code_only(tmp_path, "stale base with drifted backup") + (tmp_path / guard.SCHEMA_PATH).write_text("{}") + stale_base: Final = _commit_code_only(tmp_path, "stale base with drifted backup and schema") head: Final = _commit_code_only(tmp_path, "code change on the stale base") human: Final = _run_guard(tmp_path, stale_base, head, "litellm_fix_pricing") assert human.returncode == 0, human.stdout + human.stderr diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_schema.py similarity index 67% rename from tests/unit/test_model_prices_and_context_window.py rename to tests/unit/test_model_prices_schema.py index 49a6fb05b15..211a43d2574 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_schema.py @@ -1,5 +1,6 @@ from __future__ import annotations +import importlib.util import json import re from collections.abc import Mapping @@ -7,6 +8,7 @@ from pathlib import Path from types import MappingProxyType from typing import Final +import jsonschema import pytest import litellm @@ -15,8 +17,27 @@ from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts REPO_ROOT = Path(__file__).parents[2] +GENERATOR_PATH = REPO_ROOT / "ci_cd" / "generate_model_prices_schema.py" PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json" BACKUP_PRICES_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" +SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json" + + +def build_validator(schema: dict) -> jsonschema.Draft202012Validator: + return jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER) + + +def load_generator(): + spec = importlib.util.spec_from_file_location("generate_model_prices_schema", GENERATOR_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def committed_schema() -> dict: + return json.loads(SCHEMA_PATH.read_text()) @pytest.fixture(scope="module") @@ -24,6 +45,183 @@ def prices() -> dict: return json.loads(PRICES_PATH.read_text()) +def test_committed_schema_matches_generator_output(prices: dict, committed_schema: dict): + generator = load_generator() + regenerated = json.loads(generator.render(generator.build_schema(prices))) + assert regenerated == committed_schema, ( + "model_prices_and_context_window.schema.json is out of sync; " + "run `python ci_cd/generate_model_prices_schema.py` and commit the result" + ) + + +def test_prices_file_validates_against_committed_schema(prices: dict, committed_schema: dict): + validator = build_validator(committed_schema) + errors = [ + f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}" + for error in validator.iter_errors(prices) + ] + assert errors == [] + + +@pytest.mark.parametrize( + "entry", + [ + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": "0.01"}, + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": -1}, + {"litellm_provider": "openai", "mode": "not_a_real_mode"}, + {"mode": "chat"}, + {"litellm_provider": "openai", "deprecation_date": "June 2026"}, + {"litellm_provider": "openai", "deprecation_date": "2026-99-99"}, + {"litellm_provider": "openai", "deprecation_date": "2026-13-01"}, + {"litellm_provider": "openai", "deprecation_date": "2026-01-32"}, + {"litellm_provider": "openai", "deprecation_date": "2026-01-00"}, + {"litellm_provider": "openai", "deprecation_date": "2026-02-31"}, + {"litellm_provider": "openai", "supported_modalities": ["smell"]}, + {"litellm_provider": "openai", "supports_vision": "yes"}, + {"litellm_provider": "openai", "max_tokens": 8191.5}, + {"litellm_provider": "openai", "tiered_pricing": [{"unknown_tier_field": 1}]}, + ], + ids=[ + "cost_as_string", + "negative_cost", + "unknown_mode", + "missing_provider", + "non_iso_deprecation_date", + "impossible_month_and_day", + "month_out_of_range", + "day_out_of_range", + "day_zero", + "calendar_impossible_day", + "unknown_modality", + "boolean_flag_as_string", + "fractional_max_tokens", + "unknown_tiered_pricing_field", + ], +) +def test_schema_rejects_malformed_entries(committed_schema: dict, entry: dict): + validator = build_validator(committed_schema) + assert not validator.is_valid({"some-model": entry}) + + +def test_schema_accepts_minimal_and_unknown_optional_fields(committed_schema: dict): + validator = build_validator(committed_schema) + assert validator.is_valid({"some-model": {"litellm_provider": "openai"}}) + assert validator.is_valid({"some-model": {"litellm_provider": "openai", "brand_new_field": {"nested": True}}}) + + +def test_schema_accepts_cache_creation_cost_inside_a_pricing_tier(committed_schema: dict): + validator = build_validator(committed_schema) + entry = { + "litellm_provider": "dashscope", + "mode": "chat", + "tiered_pricing": [ + { + "range": [0, 256000], + "input_cost_per_token": 3.25e-07, + "output_cost_per_token": 1.95e-06, + "cache_creation_input_token_cost": 4.063e-07, + "cache_read_input_token_cost": 3.25e-08, + } + ], + } + assert validator.is_valid({"some-model": entry}) + + +OFF_PEAK_ENTRY: Final = MappingProxyType( + { + "litellm_provider": "openrouter", + "mode": "chat", + "input_cost_per_token": 2e-6, + "output_cost_per_token": 8e-6, + "off_peak_pricing": { + "hours_utc": "16:30-00:30", + "windows": [{"hours_utc": ["00:30-02:00"], "weekdays": [6, "Sunday", "mon", "THURS"]}], + "weekday_timezone": "Asia/Shanghai", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 4e-6, + "cache_read_input_token_cost": 1e-7, + }, + } +) + + +def test_generator_classifies_off_peak_pricing_as_a_windowed_rate_block(): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(OFF_PEAK_ENTRY)}))) + validator = build_validator(schema) + assert validator.is_valid({"some-model": dict(OFF_PEAK_ENTRY)}) + + +@pytest.mark.parametrize( + "block", + [ + {"hours_utc": "16:30-00:30", "input_cost_per_token": "1e-6"}, + {"hours_utc": "16:30-00:30", "input_cost_per_token": -1e-6}, + {"hours_utc": 1630, "input_cost_per_token": 1e-6}, + {"hours_utc": "16:30-00:30", "discount": 0.5}, + {"windows": [{"weekdays": [6]}], "input_cost_per_token": 1e-6}, + {"windows": [{"hours_utc": "00:30-02:00", "weekdays": [0]}], "input_cost_per_token": 1e-6}, + {"windows": [], "input_cost_per_token": 1e-6}, + {"input_cost_per_token": 1e-6}, + {"hours_utc": "16:30", "input_cost_per_token": 1e-6}, + {"hours_utc": "25:00-01:00", "input_cost_per_token": 1e-6}, + {"hours_utc": ["16:30-00:30", "4pm-midnight"], "input_cost_per_token": 1e-6}, + {"windows": [{"hours_utc": "00:30-02:00", "weekdays": ["Funday"]}], "input_cost_per_token": 1e-6}, + ], +) +def test_generated_off_peak_schema_rejects_malformed_blocks(block: dict): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(OFF_PEAK_ENTRY)}))) + validator = build_validator(schema) + assert not validator.is_valid({"some-model": {**OFF_PEAK_ENTRY, "off_peak_pricing": block}}) + + +TOKEN_COUNT_ENTRY: Final = MappingProxyType( + { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "token_count_framing": 78, + "token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"}, + } +) + + +@pytest.mark.parametrize( + "setting", + [ + {"token_count_framing": 0}, + {"token_count_framing": 78}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5"}}, + ], +) +def test_generator_classifies_token_count_settings(setting: dict): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) + assert build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}}) + + +@pytest.mark.parametrize( + "setting", + [ + {"token_count_framing": -1}, + {"token_count_framing": "78"}, + {"token_count_framing": 1.5}, + {"token_count_model": {"model": ""}}, + {"token_count_model": {"model": " "}}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": ""}}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": " "}}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": 5}}, + {"token_count_model": {"region": "us-east-1"}}, + {"token_count_model": "us.anthropic.claude-opus-5-5"}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "zone": "us-east-1"}}, + ], +) +def test_generated_token_count_schema_rejects_malformed_settings(setting: dict): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) + assert not build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}}) + + def find_duplicate_keys(path: Path) -> list[str]: duplicates: list[str] = [] From 2ddc8946d0c7ac74978ea12e9b0d8d214eeeed4e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:07:15 -0700 Subject: [PATCH 27/36] fix(count_tokens): complete the selected count settings prerequisites Preserve the prompt normalization required by #46777 from #46839 and classify its token_count_tokenizer key using the upstream schema definition. --- ci_cd/generate_model_prices_schema.py | 5 +++++ litellm/proxy/proxy_server.py | 4 +++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 359df8e6e32..a5a4b7f6f3b 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -259,6 +259,11 @@ def cost_schema(key: str) -> JsonSchema: def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]: return { + "token_count_tokenizer": { + "type": "string", + "description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.", + "enum": ["kimi_k3"], + }, "litellm_provider": { "type": "string", "description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4ec1cf36408..8369e79596c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14153,8 +14153,10 @@ async def count_request_tokens( model_info.get("custom_tokenizer", None), ) system_message: Final = _system_message(system) + prompt_turn: Final = 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, messages + Sequence[AllMessageValues] | None, + [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else messages), ) counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages) From e19f8fd1739f7a6d024358e8165bb6338e742416 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:08:16 -0700 Subject: [PATCH 28/36] ci(cost map): regenerate the schema for the stable line model map --- model_prices_and_context_window.schema.json | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index b99f8ee6f98..5b2d52f11ee 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -1088,24 +1088,6 @@ "minimum": 0, "description": "Tokens added once per request to a local count estimate, for the provider's chat framing." }, - "token_count_model": { - "type": "object", - "required": [ - "model" - ], - "additionalProperties": false, - "properties": { - "model": { - "type": "string", - "pattern": "\\S" - }, - "region": { - "type": "string", - "pattern": "\\S" - } - }, - "description": "Model that counts tokens for this one, and the AWS region its count API runs in." - }, "token_count_tokenizer": { "type": "string", "description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.", From d522e2e09bfe335775f2456339feb41726a285ad Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:28:17 -0700 Subject: [PATCH 29/36] fix(count_tokens): preserve legacy provider error fallback on stable Defer the breaking error policy from #46121. Adapt the #46777 count settings flow to attempt a labelled local estimate for provider failures across providers and configured count models. Keep provider errors when local counting is disabled or cannot estimate the request, and test the shared behavior across HTTP error statuses. --- litellm/proxy/proxy_server.py | 6 +- tests/unit/proxy/test_proxy_token_counter.py | 279 +++++-------------- 2 files changed, 66 insertions(+), 219 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8369e79596c..7b0607da35d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13696,10 +13696,6 @@ async def _provider_token_count( settings: TokenCountSettings, local_estimate_allowed: bool, ) -> Union["TokenCountResponse", ProviderCountFailure, None]: - """Count with the provider, using the token_count_model in its region when one is set. - - Returns None when the count should be estimated locally, or the provider's failure when the framed local - estimate stands in for it and that failure is the answer if the estimate cannot read the request.""" count_model: Final = settings.count_model counted: Final = await ( count_with(model_to_use, None) if count_model is None else count_with(count_model.model, count_model.region) @@ -13718,7 +13714,7 @@ async def _provider_token_count( "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" ) ) - if count_model is None and failure.cannot_count_model and settings.framing is not None and local_estimate_allowed: + if local_estimate_allowed: return failure raise failure.to_proxy_exception() diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 38f2181d4e7..0bfeb3fddab 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -21,6 +21,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import anyio import httpx import pytest +import respx import tiktoken from dotenv import load_dotenv from pydantic import BaseModel, ConfigDict, JsonValue @@ -40,6 +41,7 @@ from fastapi.responses import JSONResponse import litellm from litellm import Router from litellm._logging import verbose_proxy_logger +from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError @@ -362,63 +364,41 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): @pytest.mark.asyncio -async def test_internal_token_counter_anthropic_provider_detection(): - """ - Test that the internal token_counter correctly detects Anthropic providers - and handles the from_anthropic_endpoint flag appropriately - """ - - # Test with Anthropic provider - llm_router = Router( - model_list=[ - { - "model_name": "claude-test", - "litellm_params": { - "model": "anthropic/claude-3-sonnet-20240229", - "api_key": "test-key", - }, - } - ] +@pytest.mark.parametrize("status_code", (400, 401, 403, 404, 429, 500)) +@pytest.mark.parametrize( + ("model", "count_url"), + ( + ("anthropic/claude-sonnet-5-5", "https://api.anthropic.com/v1/messages/count_tokens"), + ("openai/gpt-6.1-sol", "https://api.openai.com/v1/responses/input_tokens"), + ), +) +async def test_provider_errors_return_labelled_local_estimates( + monkeypatch: pytest.MonkeyPatch, status_code: int, model: str, count_url: str +) -> None: + messages: Final = [ + {"role": "user", "content": "Count the whole request locally when the count API is unavailable."} + ] + router: Final = Router( + model_list=[{"model_name": "counted", "litellm_params": {"model": model, "api_key": "test-count-key"}}] ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "disable_token_counter", False) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") - setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - - # Test with is_direct_request=False (simulating call from Anthropic endpoint) - # 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, + with respx.mock(assert_all_called=True) as upstream: + upstream.post(count_url).mock( + return_value=httpx.Response( + status_code, json={"type": "error", "error": {"type": "api_error", "message": "Count unavailable"}} + ) ) - assert exc_info.value.type == "token_counting_error" - - # Test with non-Anthropic provider - llm_router = Router( - model_list=[ - { - "model_name": "gpt-test", - "litellm_params": { - "model": "gpt-4", - }, - } - ] - ) - - setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - - with pytest.raises(ProxyException) as missing_key: - await token_counter( - request=TokenCountRequest( - model="gpt-test", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, + counted: Final = await token_counter( + request=TokenCountRequest(model="counted", messages=messages), call_endpoint=True ) - assert missing_key.value.code == "401" - assert "No OpenAI API key is set for model gpt-test" in missing_key.value.message + + assert counted.total_tokens == litellm.token_counter(model=model.split("/", 1)[1], messages=messages) + assert counted.count_source == "local_estimate" + assert counted.error is False @pytest.mark.asyncio @@ -915,69 +895,6 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): assert exc_info.value.message == "Forbidden - Invalid credentials" -@pytest.mark.asyncio -async def test_token_counter_httpx_status_error_raises_proxy_exception(): - """ - When provider_counter.count_tokens() raises httpx.HTTPStatusError, - the token_counter endpoint should catch it and raise a ProxyException - with the upstream status code and error message. - """ - - upstream_status = 429 - upstream_message = "Rate limit exceeded" - response = httpx.Response( - status_code=upstream_status, - request=httpx.Request("POST", "https://provider.example.com/count"), - ) - http_error = httpx.HTTPStatusError( - message=upstream_message, - request=response.request, - response=response, - ) - - mock_counter = MagicMock() - mock_counter.should_use_token_counting_api.return_value = True - mock_counter.count_tokens = AsyncMock(side_effect=http_error) - - # Save originals - original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter - original_router = litellm.proxy.proxy_server.llm_router - - try: - - def mock_get_provider_token_counter(deployment, model_to_use): - return (mock_counter, "claude-4-6-sonnet", "vertex_ai") - - litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter - - mock_router = MagicMock() - mock_router.async_get_available_deployment = AsyncMock( - return_value={ - "litellm_params": { - "model": "vertex_ai/claude-4-6-sonnet", - "api_key": "fake-key", - }, - "model_info": {}, - } - ) - litellm.proxy.proxy_server.llm_router = mock_router - - with pytest.raises(ProxyException) as exc_info: - await token_counter( - request=TokenCountRequest( - model="claude-4-6-sonnet", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) - - assert exc_info.value.code == str(upstream_status) - assert upstream_message in exc_info.value.message - assert exc_info.value.type == "token_counting_error" - assert exc_info.value.param == "model" - finally: - litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter - litellm.proxy.proxy_server.llm_router = original_router @pytest.mark.asyncio @@ -1044,71 +961,6 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter -@pytest.mark.asyncio -async def test_proxy_token_counter_error_raises_when_enabled(): - """ - 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( - total_tokens=0, - request_model="bedrock/anthropic.claude-3-sonnet", - model_used="anthropic.claude-3-sonnet", - tokenizer_type="bedrock_api", - error=True, - error_message="Rate limit exceeded", - status_code=429, - ) - - # Create mock router that returns a deployment - mock_deployment = { - "litellm_params": { - "model": "bedrock/anthropic.claude-3-sonnet", - }, - "model_info": {}, - } - - mock_router = MagicMock() - mock_router.async_get_available_deployment = AsyncMock(return_value=mock_deployment) - - setattr(litellm.proxy.proxy_server, "llm_router", mock_router) - - # Save original value and function - original_disable = litellm.disable_token_counter - original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter - - try: - litellm.disable_token_counter = False - - # Create a mock counter that returns an error response - mock_counter = MagicMock(spec=BedrockTokenCounter) - mock_counter.should_use_token_counting_api.return_value = True - mock_counter.count_tokens = AsyncMock(return_value=error_response) - - # Replace the function directly - def mock_get_provider_token_counter(deployment, model_to_use): - return (mock_counter, "anthropic.claude-3-sonnet", "bedrock") - - litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter - - 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, - ) - - 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 = original_get_provider_token_counter @pytest.mark.asyncio @@ -1325,7 +1177,7 @@ _OPUS_IN_US_EAST_1: Final = TokenCountModel(model="us.anthropic.claude-opus-5-5" async def _ladder( counter: _ScriptedCounter, settings: TokenCountSettings, local_estimate_allowed: bool = True -) -> TokenCountResponse | None: +) -> TokenCountResponse | ProviderCountFailure | None: return await _provider_token_count( count_with=partial( _try_provider_token_count, @@ -1381,17 +1233,18 @@ async def test_token_count_model_counts_the_same_request_in_its_own_region_witho assert _SONNET_DEPLOYMENT["litellm_params"]["aws_region_name"] == "us-west-2" -@pytest.mark.parametrize("status_code", [401, 403, 429, 500]) +@pytest.mark.parametrize("status_code", (400, 401, 403, 404, 429, 500)) +@pytest.mark.parametrize("framing", (None, 0, 10)) @pytest.mark.asyncio -async def test_errors_other_than_cannot_count_raise_even_with_framing(status_code: int): +async def test_provider_count_errors_preserve_legacy_local_fallback(status_code: int, framing: int | None): counter: Final = _ScriptedCounter( {"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=status_code)} ) - with pytest.raises(ProxyException) as raised: - await _ladder(counter, TokenCountSettings(framing=10, count_model=None)) + result: Final = await _ladder(counter, TokenCountSettings(framing=framing, count_model=None)) - assert (raised.value.code, raised.value.message) == (str(status_code), f"scripted {status_code} for sonnet") + assert isinstance(result, ProviderCountFailure) + assert (result.status_code, result.message) == (status_code, f"scripted {status_code} for sonnet") @pytest.mark.parametrize("status_code", [400, 403, 404, 500]) @@ -1401,38 +1254,36 @@ async def test_token_count_model_failure_names_its_region_so_the_user_can_enable {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} ) - with pytest.raises(ProxyException) as raised: - await _ladder(counter, TokenCountSettings(framing=None, count_model=_OPUS_IN_US_EAST_1)) + result: Final = await _ladder(counter, TokenCountSettings(framing=None, count_model=_OPUS_IN_US_EAST_1)) - assert raised.value.code == str(status_code) - assert raised.value.message.startswith(f"scripted {status_code} for opus") - assert "us.anthropic.claude-opus-5-5 in us-east-1" in raised.value.message + assert isinstance(result, ProviderCountFailure) + assert result.status_code == status_code + assert result.message.startswith(f"scripted {status_code} for opus") + assert "us.anthropic.claude-opus-5-5 in us-east-1" in result.message assert len(counter.calls) == 1 @pytest.mark.asyncio -async def test_token_count_model_without_a_region_answers_its_error_unchanged(): +async def test_token_count_model_without_a_region_preserves_its_error_for_a_failed_local_estimate(): counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=404)}) - with pytest.raises(ProxyException) as raised: - await _ladder( - counter, - TokenCountSettings( - framing=None, count_model=TokenCountModel(model="us.anthropic.claude-opus-5-5", region=None) - ), - ) + result: Final = await _ladder( + counter, + TokenCountSettings(framing=None, count_model=TokenCountModel(model="us.anthropic.claude-opus-5-5", region=None)), + ) - assert (raised.value.code, raised.value.message) == ("404", "scripted 404 for opus") + assert isinstance(result, ProviderCountFailure) + assert (result.status_code, result.message) == (404, "scripted 404 for opus") @pytest.mark.asyncio -async def test_uncountable_model_without_framing_raises(): +async def test_uncountable_model_without_framing_preserves_legacy_local_fallback(): counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=404)}) - with pytest.raises(ProxyException) as raised: - await _ladder(counter, TokenCountSettings(framing=None, count_model=None)) + result: Final = await _ladder(counter, TokenCountSettings(framing=None, count_model=None)) - assert raised.value.code == "404" + assert isinstance(result, ProviderCountFailure) + assert (result.status_code, result.message) == (404, "scripted 404 for sonnet") @pytest.mark.asyncio @@ -1447,16 +1298,16 @@ async def test_uncountable_model_with_framing_falls_through_to_the_local_estimat @pytest.mark.parametrize("status_code", [400, 404]) @pytest.mark.asyncio -async def test_token_count_model_error_stands_even_with_framing(status_code: int): +async def test_token_count_model_failure_allows_a_framed_local_estimate(status_code: int): counter: Final = _ScriptedCounter( {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} ) - with pytest.raises(ProxyException) as raised: - await _ladder(counter, TokenCountSettings(framing=78, count_model=_OPUS_IN_US_EAST_1)) + result: Final = await _ladder(counter, TokenCountSettings(framing=78, count_model=_OPUS_IN_US_EAST_1)) - assert raised.value.code == str(status_code) - assert "us.anthropic.claude-opus-5-5 in us-east-1" in raised.value.message + assert isinstance(result, ProviderCountFailure) + assert result.status_code == status_code + assert "us.anthropic.claude-opus-5-5 in us-east-1" in result.message @pytest.mark.asyncio @@ -1470,7 +1321,7 @@ async def test_uncountable_model_with_framing_raises_when_local_counting_is_disa @pytest.mark.asyncio -async def test_request_the_counter_rejects_before_calling_the_provider_answers_400_without_framing(): +async def test_rejected_provider_conversion_still_allows_a_local_estimate(): counter: Final = _ScriptedCounter( { "us.anthropic.claude-sonnet-5-5": _provider_reply( @@ -1479,10 +1330,10 @@ async def test_request_the_counter_rejects_before_calling_the_provider_answers_4 } ) - with pytest.raises(ProxyException) as raised: - await _ladder(counter, TokenCountSettings(framing=100, count_model=None)) + result: Final = await _ladder(counter, TokenCountSettings(framing=100, count_model=None)) - assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for sonnet") + assert isinstance(result, ProviderCountFailure) + assert (result.status_code, result.message) == (400, "scripted 400 for sonnet") @pytest.mark.parametrize("model", ["openai/gpt-5.4-mini", "gemini/gemini-3.8-flash"]) From e9e7e32142dcebc7080fe0e27583ca443bdb1c25 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:30:57 -0700 Subject: [PATCH 30/36] test(count_tokens): make strict error envelope coverage explicit --- .../providers/test_count_tokens_error_shapes_wire.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py index 200f5e81673..74581e44310 100644 --- a/tests/integration/providers/test_count_tokens_error_shapes_wire.py +++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py @@ -1,5 +1,6 @@ import json from collections.abc import Iterator +from pathlib import Path from typing import Final import anthropic @@ -29,13 +30,22 @@ _PROVIDER_KEY_VARIABLES: Final = ( @pytest.fixture(scope="module") -def keyless_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: +def provider_error_config(tmp_path_factory: pytest.TempPathFactory) -> Path: + config: Final = tmp_path_factory.mktemp("count-tokens-provider-errors") / "config.yaml" + base: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" + config.write_text(base.read_text() + "\nlitellm_settings:\n disable_token_counter: true\n") + return config + + +@pytest.fixture(scope="module") +def keyless_proxy(tmp_path_factory: pytest.TempPathFactory, provider_error_config: Path) -> Iterator[OwnedProxy]: with ( gateway_from_environment() as gateway, owned_proxy_process( gateway, tmp_path_factory.mktemp("count-tokens-error-shapes"), {}, + config=provider_error_config, remove_environment=_PROVIDER_KEY_VARIABLES, workers=1, ) as owned, From 498095ee1b887ab16020b70ef0dd60f63351d58c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:38:30 -0700 Subject: [PATCH 31/36] fix(count_tokens): retain legacy translated content fallback --- litellm/proxy/proxy_server.py | 12 +--- tests/unit/proxy/test_proxy_token_counter.py | 66 ++++++++++---------- 2 files changed, 35 insertions(+), 43 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7b0607da35d..63a920628b6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -330,7 +330,6 @@ from litellm.litellm_core_utils.core_helpers import ( as_str_mapping, drop_params_flag, get_litellm_metadata_from_kwargs, - iter_items, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -14002,14 +14001,6 @@ def _drops_params(deployment: Mapping[str, object] | None) -> bool: return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) -_ESTIMABLE_GOOGLE_PART_KEYS: Final = frozenset({"text", "functionCall", "functionResponse", "thoughtSignature"}) - - -def _estimable_contents(contents: Sequence[Mapping[str, object]]) -> bool: - parts: Final = chain.from_iterable(iter_items(content.get("parts") or ()) for content in contents) - return all(set(as_str_mapping(part) or ()) <= _ESTIMABLE_GOOGLE_PART_KEYS for part in parts) - - def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -14105,7 +14096,6 @@ async def count_request_tokens( code=400, ) prompt_only: Final = bool(prompt) and conversation_empty - 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 ladder: Final = ( None @@ -14126,7 +14116,7 @@ async def count_request_tokens( ), model_to_use=model_to_use, settings=settings, - local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable, + local_estimate_allowed=litellm.disable_token_counter is not True, ) ) if isinstance(ladder, TokenCountResponse): diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 0bfeb3fddab..995089105c4 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -401,6 +401,40 @@ async def test_provider_errors_return_labelled_local_estimates( assert counted.error is False +@pytest.mark.asyncio +@pytest.mark.parametrize( + "part", + ( + {"fileData": {"mimeType": "image/png", "fileUri": "gs://bucket/image.png"}}, + {"inlineData": {"mimeType": "image/png", "data": "aW1hZ2U="}}, + {"text": "Reasoning about the request", "thought": True}, + ), +) +async def test_google_provider_errors_preserve_legacy_translated_message_estimates( + monkeypatch: pytest.MonkeyPatch, part: dict[str, JsonValue] +) -> None: + model: Final = "openai/gpt-6.1-sol" + request: Final = _google_count_tokens_request( + "counted", {"contents": [{"role": "user", "parts": [{"text": "Describe this"}, part]}]} + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "counted", "litellm_params": {"model": model, "api_key": "test-key"}}]), + ) + monkeypatch.setattr(litellm, "disable_token_counter", False) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + with respx.mock(assert_all_called=True) as upstream: + upstream.post("https://api.openai.com/v1/responses/input_tokens").mock( + return_value=httpx.Response(400, json={"error": {"message": "Count unavailable"}}) + ) + counted: Final = await token_counter(request=request, call_endpoint=True) + assert counted.total_tokens == litellm.token_counter(model="gpt-6.1-sol", messages=request.messages) + assert counted.count_source == "local_estimate" + assert counted.error is False + + @pytest.mark.asyncio async def test_anthropic_endpoint_error_handling(): """ @@ -1410,38 +1444,6 @@ async def test_framing_keeps_the_400_for_blocks_the_provider_counter_cannot_conv }, id="array_schema_without_items", ), - pytest.param( - {"messages": [], "contents": [{"role": "user", "parts": [{"text": "What is new?"}]}]}, - id="google_contents", - ), - pytest.param( - { - "messages": [{"role": "user", "content": "Describe this"}], - "contents": [ - { - "role": "user", - "parts": [ - {"text": "Describe this"}, - {"fileData": {"mimeType": "image/png", "fileUri": "gs://b/c.png"}}, - ], - } - ], - }, - id="google_contents_with_media_the_translation_dropped", - ), - pytest.param( - { - "messages": [ - {"role": "user", "content": "What is new?"}, - {"role": "assistant", "content": "Weighing the release notes"}, - ], - "contents": [ - {"role": "user", "parts": [{"text": "What is new?"}]}, - {"role": "model", "parts": [{"text": "Weighing the release notes", "thought": True}]}, - ], - }, - id="google_contents_with_a_thought_the_translation_counts_as_text", - ), ], ) async def test_framed_estimate_that_cannot_read_the_body_answers_the_provider_error( From c881f7cfc169ca187bf29a2a3f2a68f139cae0b6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:10:51 -0700 Subject: [PATCH 32/36] fix(count_tokens): bound PDF parsing and preserve legacy fallback --- litellm/litellm_core_utils/token_counter.py | 110 ++++++++++++-- pyproject.toml | 2 +- .../litellm_core_utils/test_token_counter.py | 141 ++++++++++++++++-- uv.lock | 2 +- 4 files changed, 230 insertions(+), 25 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index c62233f49e2..07c429e5607 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -4,8 +4,9 @@ import base64 import io import math import struct -from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence -from typing import Final, Literal, cast +from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence +from itertools import accumulate +from typing import TYPE_CHECKING, Final, Literal, cast import anyio import anyio.lowlevel @@ -13,6 +14,10 @@ import httpx import tiktoken from typing_extensions import ParamSpec, TypeVar +if TYPE_CHECKING: + from pypdf import PageObject, PdfReader + from pypdf.generic import IndirectObject + import litellm from litellm import verbose_logger from litellm._lazy_imports import _get_default_encoding @@ -784,23 +789,102 @@ def _anthropic_rendered_page_image_tokens(width: float, height: float) -> int: return math.ceil(min(float(ANTHROPIC_IMAGE_MAX_PIXELS), area_at_max_edge) / ANTHROPIC_IMAGE_PIXELS_PER_TOKEN) +_PDF_MAX_INPUT_BYTES: Final = 32 * 1024 * 1024 +_PDF_MAX_PAGES: Final = 100 +_PDF_MAX_OBJECTS: Final = 2048 +_PDF_MAX_STREAM_BYTES: Final = 1024 * 1024 +_PDF_MAX_TOTAL_STREAM_BYTES: Final = 4 * 1024 * 1024 + + +class _PdfTextBudget: + def __init__(self) -> None: + self.operations_left: int = 100_000 + self.external_objects_left: int = 16 + self.exceeded: bool = False + + def visit(self, operator: bytes, _operands: Sequence[object], _cm: Sequence[float], _tm: Sequence[float]) -> None: + from pypdf.errors import LimitReachedError + + self.operations_left -= 1 + self.external_objects_left -= int(operator == b"Do") + if self.operations_left < 0 or self.external_objects_left < 0: + self.exceeded = True + raise LimitReachedError("PDF text extraction exceeds the counting budget") + + +def _pdf_page_tokens(page: "PageObject", count_function: TokenCounterFunction) -> int | None: + budget: Final = _PdfTextBudget() + text: Final = page.extract_text(visitor_operand_before=budget.visit) + if budget.exceeded: + return None + return count_function(text or "") + _anthropic_rendered_page_image_tokens( + float(page.mediabox.width), float(page.mediabox.height) + ) + + +def _pdf_references(reader: "PdfReader") -> Iterator["IndirectObject"]: + from pypdf.generic import IndirectObject + + for generation, identifiers in reader.xref.items(): + for identifier in identifiers: + if identifier: + yield IndirectObject(identifier, generation, reader) + for identifier in reader.xref_objStm: + yield IndirectObject(identifier, 0, reader) + + +def _pdf_stream_sizes(reader: "PdfReader") -> Iterator[int]: + from pypdf.generic import StreamObject + + for reference in _pdf_references(reader): + value: Final = reader.get_object(reference) + if isinstance(value, StreamObject): + yield len(value.get_data()) + + def _count_inline_pdf_tokens(data_url: str, count_function: TokenCounterFunction) -> int | None: if not data_url.startswith(PDF_DATA_URL_PREFIX): return None - try: - from pypdf import PdfReader - except ImportError: - verbose_logger.debug("pypdf is not installed, so the PDF document is priced like one image") + encoded_length: Final = len(data_url) - len(PDF_DATA_URL_PREFIX) + if encoded_length > 4 * ((_PDF_MAX_INPUT_BYTES + 2) // 3): return None try: - reader: Final = PdfReader(io.BytesIO(base64.b64decode(data_url[len(PDF_DATA_URL_PREFIX) :]))) - return sum( - count_function(page.extract_text() or "") - + _anthropic_rendered_page_image_tokens(float(page.mediabox.width), float(page.mediabox.height)) - for page in reader.pages - ) + from pypdf import PdfReader, apply_configuration + except ImportError: + verbose_logger.debug("A bounded pypdf parser is unavailable, so the PDF document is priced like one image") + return None + try: + decoded: Final = base64.b64decode(data_url[len(PDF_DATA_URL_PREFIX) :], validate=True) + if len(decoded) > _PDF_MAX_INPUT_BYTES: + return None + with apply_configuration( + maximum_declared_stream_length=_PDF_MAX_INPUT_BYTES, + array_based_stream_maximum_output_length=_PDF_MAX_STREAM_BYTES, + zlib_maximum_output_length=_PDF_MAX_STREAM_BYTES, + zlib_maximum_recovery_input_length=_PDF_MAX_STREAM_BYTES, + lzw_maximum_output_length=_PDF_MAX_STREAM_BYTES, + run_length_maximum_output_length=_PDF_MAX_STREAM_BYTES, + jbig2_maximum_output_length=_PDF_MAX_STREAM_BYTES, + image_maximum_buffer_size=_PDF_MAX_STREAM_BYTES, + page_tree_maximum_entries=2 * _PDF_MAX_PAGES + 1, + page_tree_maximum_depth=16, + xform_maximum_invocations_per_extraction=16, + jbig2dec_binary=None, + ): + reader: Final = PdfReader(io.BytesIO(decoded), strict=True) + object_count: Final = sum(len(identifiers) for identifiers in reader.xref.values()) + len( + reader.xref_objStm + ) + if object_count > _PDF_MAX_OBJECTS or len(reader.pages) > _PDF_MAX_PAGES: + return None + if not all(size <= _PDF_MAX_TOTAL_STREAM_BYTES for size in accumulate(_pdf_stream_sizes(reader))): + return None + page_counts: Final = tuple(_pdf_page_tokens(page, count_function) for page in reader.pages) + return None if None in page_counts else sum(count for count in page_counts if count is not None) except Exception as e: - verbose_logger.debug("Could not read the PDF document's pages (%s), so it is priced like one image", e) + verbose_logger.debug( + "Could not read the PDF document within counting limits (%s), so it is priced like one image", e + ) return None diff --git a/pyproject.toml b/pyproject.toml index 2816de2316c..89c98f8b7f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -181,7 +181,7 @@ proxy-runtime = [ "mangum>=0.17.0,<1.0", "azure-ai-contentsafety>=1.0.0,<2.0", "azure-storage-file-datalake>=12.20.0,<13.0", - "pypdf>=6.16.1,<7.0", + "pypdf>=6.19.0,<7.0", "llm-sandbox>=0.3.39,<1.0", "detect-secrets>=1.5.0,<2.0", ] diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 1285a291650..81e6062ae93 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -3,12 +3,14 @@ import asyncio import base64 import importlib +import io import json import os import subprocess import sys import threading import time +import zlib from collections.abc import Mapping from concurrent.futures import Future, wait from itertools import accumulate, chain @@ -1448,16 +1450,25 @@ def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): assert _count_user_content([prompt, by_id]) == _count_user_content([prompt]) named = {"type": "file", "file": {"file_id": "file-abc123", "filename": "report.pdf"}} - assert _count_user_content([prompt, named]) == _count_user_content( - [prompt, {"type": "text", "text": "report.pdf"}] - ) + assert _count_user_content([prompt, named]) == _count_user_content([prompt, {"type": "text", "text": "report.pdf"}]) -def _pdf_base64(pages: tuple[str, ...], width: int = 612, height: int = 792) -> str: +def _pdf_base64( + pages: tuple[str, ...], + width: int = 612, + height: int = 792, + *, + padding: int = 0, + compressed: bool = False, + unused_objects: int = 0, + unused_bytes: int = 0, +) -> str: def page_objects(index: int, text: str) -> tuple[bytes, bytes]: escaped: Final = text.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)") - stream: Final = f"BT /F1 12 Tf 72 720 Td ({escaped}) Tj ET".encode("latin-1") - content: Final = f"<< /Length {len(stream)} >>\nstream\n".encode() + stream + b"\nendstream" + raw_stream: Final = b" " * padding + f"BT /F1 12 Tf 72 720 Td ({escaped}) Tj ET".encode("latin-1") + stream: Final = zlib.compress(raw_stream) if compressed else raw_stream + filter_field: Final = " /Filter /FlateDecode" if compressed else "" + content: Final = f"<< /Length {len(stream)}{filter_field} >>\nstream\n".encode() + stream + b"\nendstream" page: Final = ( f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {width} {height}] " f"/Resources << /Font << /F1 3 0 R >> >> /Contents {4 + 2 * index} 0 R >>" @@ -1470,6 +1481,8 @@ def _pdf_base64(pages: tuple[str, ...], width: int = 612, height: int = 792) -> f"<< /Type /Pages /Count {len(pages)} /Kids [ {kids} ] >>".encode(), b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>", *chain.from_iterable(page_objects(index, text) for index, text in enumerate(pages)), + *(b"0" for _ in range(unused_objects)), + *((b"(" + b" " * unused_bytes + b")",) if unused_bytes else ()), ) header: Final = b"%PDF-1.4\n" objects: Final = tuple( @@ -1490,6 +1503,107 @@ def _pdf_source(pdf_base64: str) -> dict[str, str]: return {"type": "base64", "media_type": "application/pdf", "data": pdf_base64} +@pytest.mark.parametrize("case", ("pages", "compressed_stream", "total_streams", "objects", "input_bytes")) +def test_pdf_counting_over_budget_preserves_legacy_image_estimate(case: str): + from litellm.litellm_core_utils.token_counter import ( + _PDF_MAX_INPUT_BYTES, + _PDF_MAX_OBJECTS, + _PDF_MAX_PAGES, + _PDF_MAX_STREAM_BYTES, + _PDF_MAX_TOTAL_STREAM_BYTES, + ) + + pdf: Final = { + "pages": lambda: _pdf_base64(("Count this",) * (_PDF_MAX_PAGES + 1)), + "compressed_stream": lambda: _pdf_base64(("Count this",), padding=_PDF_MAX_STREAM_BYTES, compressed=True), + "total_streams": lambda: _pdf_base64( + ("Count this",) * (_PDF_MAX_TOTAL_STREAM_BYTES // _PDF_MAX_STREAM_BYTES + 1), + padding=_PDF_MAX_STREAM_BYTES - 100, + compressed=True, + ), + "objects": lambda: _pdf_base64(("Count this",), unused_objects=_PDF_MAX_OBJECTS), + "input_bytes": lambda: _pdf_base64(("Count this",), unused_bytes=_PDF_MAX_INPUT_BYTES), + }[case]() + source: Final = _pdf_source(pdf) + assert _count_user_content([{"type": "document", "source": source}]) == _count_user_content( + [{"type": "image", "source": source}] + ) + + +def test_compressed_pdf_within_budget_keeps_text_and_page_counts(): + pages: Final = ("Count the first page", "Count the second page") + compressed: Final = _pdf_source(_pdf_base64(pages, compressed=True)) + uncompressed: Final = _pdf_source(_pdf_base64(pages)) + assert _count_user_content([{"type": "document", "source": compressed}]) == _count_user_content( + [{"type": "document", "source": uncompressed}] + ) + + +@pytest.mark.parametrize("nested", (False, True)) +def test_pdf_form_budget_never_returns_a_partial_count(nested: bool): + from pypdf import PdfWriter + from pypdf.generic import ArrayObject, DecodedStreamObject, DictionaryObject, NameObject, NumberObject + + writer: Final = PdfWriter() + page: Final = writer.add_blank_page(width=612, height=792) + font: Final = DictionaryObject( + { + NameObject("/Type"): NameObject("/Font"), + NameObject("/Subtype"): NameObject("/Type1"), + NameObject("/BaseFont"): NameObject("/Helvetica"), + } + ) + form: Final = DecodedStreamObject() + form.set_data(b"BT /F1 12 Tf 72 720 Td (Every form must be counted) Tj ET") + form.update( + { + NameObject("/Type"): NameObject("/XObject"), + NameObject("/Subtype"): NameObject("/Form"), + NameObject("/BBox"): ArrayObject([NumberObject(0), NumberObject(0), NumberObject(612), NumberObject(792)]), + NameObject("/Resources"): DictionaryObject( + {NameObject("/Font"): DictionaryObject({NameObject("/F1"): font})} + ), + } + ) + resources: Final = DictionaryObject( + {NameObject("/XObject"): DictionaryObject({NameObject("/Fm"): writer._add_object(form)})} + ) + repeated: Final = DecodedStreamObject() + repeated.set_data(b"/Fm Do\n" * 17) + repeated.update( + { + NameObject("/Subtype"): NameObject("/Form"), + NameObject("/BBox"): form["/BBox"], + NameObject("/Resources"): resources, + } + ) + outer: Final = DecodedStreamObject() + outer.set_data(b"/Fm Do") + page[NameObject("/Resources")] = ( + DictionaryObject({NameObject("/XObject"): DictionaryObject({NameObject("/Fm"): writer._add_object(repeated)})}) + if nested + else resources + ) + page[NameObject("/Contents")] = writer._add_object(outer if nested else repeated) + output: Final = io.BytesIO() + writer.write(output) + source: Final = _pdf_source(base64.b64encode(output.getvalue()).decode()) + assert _count_user_content([{"type": "document", "source": source}]) == _count_user_content( + [{"type": "image", "source": source}] + ) + + +def test_pdf_counting_restores_the_callers_parser_limits(): + from pypdf import get_configuration + + original: Final = get_configuration() + source: Final = _pdf_source(_pdf_base64(("Count this",))) + assert _count_user_content([{"type": "document", "source": source}]) > _count_user_content( + [{"type": "image", "source": source}] + ) + assert get_configuration() == original + + def test_base64_pdf_document_counts_every_page_text_and_rendering(): """A base64 PDF is read page by page: each page costs its text plus the image Anthropic renders it to. @@ -1500,13 +1614,18 @@ def test_base64_pdf_document_counts_every_page_text_and_rendering(): first: Final = "Revenue grew eleven percent while churn fell to two percent." second: Final = "Headcount is flat and the office lease was renewed for three years." base: Final = _count_user_content([prompt]) - blank_page: Final = _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64(("",)))}]) - base + blank_page: Final = ( + _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64(("",)))}]) - base + ) assert blank_page > 0 assert _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64((first, second)))}]) == ( - _count_user_content([prompt, {"type": "text", "text": first}, {"type": "text", "text": second}]) + 2 * blank_page + _count_user_content([prompt, {"type": "text", "text": first}, {"type": "text", "text": second}]) + + 2 * blank_page + ) + one_page: Final = ( + _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64((first,)))}]) - base ) - one_page: Final = _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64((first,)))}]) - base twelve_pages: Final = ( _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64((first,) * 12))}]) - base ) @@ -1514,7 +1633,9 @@ def test_base64_pdf_document_counts_every_page_text_and_rendering(): assert twelve_pages == 12 * one_page -@pytest.mark.parametrize("fields", [{}, {"title": "Q3 board packet", "context": "Shared by finance"}], ids=["bare", "described"]) +@pytest.mark.parametrize( + "fields", [{}, {"title": "Q3 board packet", "context": "Shared by finance"}], ids=["bare", "described"] +) def test_pdf_page_rendering_cost_follows_anthropic_image_scaling(fields: dict[str, str]): """Anthropic rasterizes each PDF page within its image limits (1568 px long edge, 1.15 MP) and bills width * height / 750 tokens for it: https://platform.claude.com/docs/en/build-with-claude/pdf-support and diff --git a/uv.lock b/uv.lock index fb9d8bbd9ac..9cc1507ecf1 100644 --- a/uv.lock +++ b/uv.lock @@ -4825,7 +4825,7 @@ requires-dist = [ { name = "pydantic-settings", specifier = ">=2.14.1,<3.0" }, { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" }, { name = "pynacl", marker = "extra == 'proxy'", specifier = ">=1.6.2,<2.0" }, - { name = "pypdf", marker = "extra == 'proxy-runtime'", specifier = ">=6.16.1,<7.0" }, + { name = "pypdf", marker = "extra == 'proxy-runtime'", specifier = ">=6.19.0,<7.0" }, { name = "pyroscope-io", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.8.16,<1.0" }, { name = "python-dotenv", specifier = ">=1.0.0,<2.0" }, { name = "python-multipart", marker = "extra == 'proxy'", specifier = ">=0.0.27,<1.0" }, From 704a15720a8e7aef22a31b577954b470eb23eab1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:14:43 -0700 Subject: [PATCH 33/36] test(count_tokens): type PDF budget regression cases --- tests/unit/litellm_core_utils/test_token_counter.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 81e6062ae93..4b260c387b7 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -1504,7 +1504,7 @@ def _pdf_source(pdf_base64: str) -> dict[str, str]: @pytest.mark.parametrize("case", ("pages", "compressed_stream", "total_streams", "objects", "input_bytes")) -def test_pdf_counting_over_budget_preserves_legacy_image_estimate(case: str): +def test_pdf_counting_over_budget_preserves_legacy_image_estimate(case: str) -> None: from litellm.litellm_core_utils.token_counter import ( _PDF_MAX_INPUT_BYTES, _PDF_MAX_OBJECTS, @@ -1530,7 +1530,7 @@ def test_pdf_counting_over_budget_preserves_legacy_image_estimate(case: str): ) -def test_compressed_pdf_within_budget_keeps_text_and_page_counts(): +def test_compressed_pdf_within_budget_keeps_text_and_page_counts() -> None: pages: Final = ("Count the first page", "Count the second page") compressed: Final = _pdf_source(_pdf_base64(pages, compressed=True)) uncompressed: Final = _pdf_source(_pdf_base64(pages)) @@ -1540,7 +1540,7 @@ def test_compressed_pdf_within_budget_keeps_text_and_page_counts(): @pytest.mark.parametrize("nested", (False, True)) -def test_pdf_form_budget_never_returns_a_partial_count(nested: bool): +def test_pdf_form_budget_never_returns_a_partial_count(nested: bool) -> None: from pypdf import PdfWriter from pypdf.generic import ArrayObject, DecodedStreamObject, DictionaryObject, NameObject, NumberObject @@ -1593,7 +1593,7 @@ def test_pdf_form_budget_never_returns_a_partial_count(nested: bool): ) -def test_pdf_counting_restores_the_callers_parser_limits(): +def test_pdf_counting_restores_the_callers_parser_limits() -> None: from pypdf import get_configuration original: Final = get_configuration() From 3f3fb2ac60073bae898002a519a39133dcc4c8a6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:20:52 -0700 Subject: [PATCH 34/36] fix(count_tokens): omit local PDF parsing from stable backport --- litellm/constants.py | 4 - litellm/litellm_core_utils/token_counter.py | 148 +-------- pyproject.toml | 2 +- tests/integration/_support/pdf_document.py | 149 --------- .../test_pdf_document_count_tokens_chaos.py | 226 -------------- .../test_pdf_document_count_tokens_wire.py | 273 ---------------- ...df_document_pre_call_checks_owned_proxy.py | 292 ------------------ .../litellm_core_utils/test_token_counter.py | 274 +++------------- uv.lock | 2 +- 9 files changed, 59 insertions(+), 1311 deletions(-) delete mode 100644 tests/integration/_support/pdf_document.py delete mode 100644 tests/integration/providers/test_pdf_document_count_tokens_chaos.py delete mode 100644 tests/integration/providers/test_pdf_document_count_tokens_wire.py delete mode 100644 tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py diff --git a/litellm/constants.py b/litellm/constants.py index b37a648835c..3b14d99bfc3 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -649,10 +649,6 @@ EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE: Final = float( ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv("ANTHROPIC_TOKEN_COUNTING_BETA_VERSION", "token-counting-2024-11-01") ANTHROPIC_SKILLS_API_BETA_VERSION: Final = "skills-2025-10-02" ANTHROPIC_BATCHES_ROUTE: Final = "/v1/messages/batches" -ANTHROPIC_IMAGE_MAX_LONG_EDGE_PX: Final = 1568 -ANTHROPIC_IMAGE_MAX_PIXELS: Final = 1_150_000 -ANTHROPIC_IMAGE_PIXELS_PER_TOKEN: Final = 750 -PDF_DATA_URL_PREFIX: Final = "data:application/pdf;base64," VERTEX_BATCH_PREDICTION_JOBS_ROUTE: Final = "batchPredictionJobs" ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES: Final = { "low": 1, diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 07c429e5607..5d7956059e4 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -2,11 +2,9 @@ ## Helper utilities for token counting import base64 import io -import math import struct -from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence -from itertools import accumulate -from typing import TYPE_CHECKING, Final, Literal, cast +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence +from typing import Final, Literal, cast import anyio import anyio.lowlevel @@ -14,17 +12,10 @@ import httpx import tiktoken from typing_extensions import ParamSpec, TypeVar -if TYPE_CHECKING: - from pypdf import PageObject, PdfReader - from pypdf.generic import IndirectObject - import litellm from litellm import verbose_logger from litellm._lazy_imports import _get_default_encoding from litellm.constants import ( - ANTHROPIC_IMAGE_MAX_LONG_EDGE_PX, - ANTHROPIC_IMAGE_MAX_PIXELS, - ANTHROPIC_IMAGE_PIXELS_PER_TOKEN, DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_TOKEN_COUNT, DEFAULT_IMAGE_WIDTH, @@ -33,7 +24,6 @@ from litellm.constants import ( MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES, MAX_TILE_HEIGHT, MAX_TILE_WIDTH, - PDF_DATA_URL_PREFIX, TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS, TOKEN_COUNTER_MAX_CONCURRENT_COUNTS, TOKEN_COUNTER_MAX_EXACT_CHARS, @@ -781,128 +771,6 @@ def _anthropic_image_source_data( return "" -def _anthropic_rendered_page_image_tokens(width: float, height: float) -> int: - """A PDF page is rasterized within Anthropic's image limits, then billed by its pixel area.""" - if width <= 0 or height <= 0: - return 0 - area_at_max_edge: Final = ANTHROPIC_IMAGE_MAX_LONG_EDGE_PX**2 * min(width, height) / max(width, height) - return math.ceil(min(float(ANTHROPIC_IMAGE_MAX_PIXELS), area_at_max_edge) / ANTHROPIC_IMAGE_PIXELS_PER_TOKEN) - - -_PDF_MAX_INPUT_BYTES: Final = 32 * 1024 * 1024 -_PDF_MAX_PAGES: Final = 100 -_PDF_MAX_OBJECTS: Final = 2048 -_PDF_MAX_STREAM_BYTES: Final = 1024 * 1024 -_PDF_MAX_TOTAL_STREAM_BYTES: Final = 4 * 1024 * 1024 - - -class _PdfTextBudget: - def __init__(self) -> None: - self.operations_left: int = 100_000 - self.external_objects_left: int = 16 - self.exceeded: bool = False - - def visit(self, operator: bytes, _operands: Sequence[object], _cm: Sequence[float], _tm: Sequence[float]) -> None: - from pypdf.errors import LimitReachedError - - self.operations_left -= 1 - self.external_objects_left -= int(operator == b"Do") - if self.operations_left < 0 or self.external_objects_left < 0: - self.exceeded = True - raise LimitReachedError("PDF text extraction exceeds the counting budget") - - -def _pdf_page_tokens(page: "PageObject", count_function: TokenCounterFunction) -> int | None: - budget: Final = _PdfTextBudget() - text: Final = page.extract_text(visitor_operand_before=budget.visit) - if budget.exceeded: - return None - return count_function(text or "") + _anthropic_rendered_page_image_tokens( - float(page.mediabox.width), float(page.mediabox.height) - ) - - -def _pdf_references(reader: "PdfReader") -> Iterator["IndirectObject"]: - from pypdf.generic import IndirectObject - - for generation, identifiers in reader.xref.items(): - for identifier in identifiers: - if identifier: - yield IndirectObject(identifier, generation, reader) - for identifier in reader.xref_objStm: - yield IndirectObject(identifier, 0, reader) - - -def _pdf_stream_sizes(reader: "PdfReader") -> Iterator[int]: - from pypdf.generic import StreamObject - - for reference in _pdf_references(reader): - value: Final = reader.get_object(reference) - if isinstance(value, StreamObject): - yield len(value.get_data()) - - -def _count_inline_pdf_tokens(data_url: str, count_function: TokenCounterFunction) -> int | None: - if not data_url.startswith(PDF_DATA_URL_PREFIX): - return None - encoded_length: Final = len(data_url) - len(PDF_DATA_URL_PREFIX) - if encoded_length > 4 * ((_PDF_MAX_INPUT_BYTES + 2) // 3): - return None - try: - from pypdf import PdfReader, apply_configuration - except ImportError: - verbose_logger.debug("A bounded pypdf parser is unavailable, so the PDF document is priced like one image") - return None - try: - decoded: Final = base64.b64decode(data_url[len(PDF_DATA_URL_PREFIX) :], validate=True) - if len(decoded) > _PDF_MAX_INPUT_BYTES: - return None - with apply_configuration( - maximum_declared_stream_length=_PDF_MAX_INPUT_BYTES, - array_based_stream_maximum_output_length=_PDF_MAX_STREAM_BYTES, - zlib_maximum_output_length=_PDF_MAX_STREAM_BYTES, - zlib_maximum_recovery_input_length=_PDF_MAX_STREAM_BYTES, - lzw_maximum_output_length=_PDF_MAX_STREAM_BYTES, - run_length_maximum_output_length=_PDF_MAX_STREAM_BYTES, - jbig2_maximum_output_length=_PDF_MAX_STREAM_BYTES, - image_maximum_buffer_size=_PDF_MAX_STREAM_BYTES, - page_tree_maximum_entries=2 * _PDF_MAX_PAGES + 1, - page_tree_maximum_depth=16, - xform_maximum_invocations_per_extraction=16, - jbig2dec_binary=None, - ): - reader: Final = PdfReader(io.BytesIO(decoded), strict=True) - object_count: Final = sum(len(identifiers) for identifiers in reader.xref.values()) + len( - reader.xref_objStm - ) - if object_count > _PDF_MAX_OBJECTS or len(reader.pages) > _PDF_MAX_PAGES: - return None - if not all(size <= _PDF_MAX_TOTAL_STREAM_BYTES for size in accumulate(_pdf_stream_sizes(reader))): - return None - page_counts: Final = tuple(_pdf_page_tokens(page, count_function) for page in reader.pages) - return None if None in page_counts else sum(count for count in page_counts if count is not None) - except Exception as e: - verbose_logger.debug( - "Could not read the PDF document within counting limits (%s), so it is priced like one image", e - ) - return None - - -def _count_opaque_document_tokens( - data_url: str, - count_function: TokenCounterFunction, - use_default_image_token_count: bool, -) -> int: - pdf_tokens: Final = _count_inline_pdf_tokens(data_url, count_function) - if pdf_tokens is not None: - return pdf_tokens - return calculate_img_tokens( - data=data_url, - mode="auto", - use_default_image_token_count=use_default_image_token_count, - ) - - def _count_document_tokens( document: ChatCompletionDocumentObject | AnthropicMessagesDocumentParam, count_function: TokenCounterFunction, @@ -922,8 +790,10 @@ def _count_document_tokens( return metadata_tokens + _count_content_list( count_function, content, use_default_image_token_count, default_token_count ) - return metadata_tokens + _count_opaque_document_tokens( - _anthropic_image_source_data(source), count_function, use_default_image_token_count + return metadata_tokens + calculate_img_tokens( + data=_anthropic_image_source_data(source), + mode="auto", + use_default_image_token_count=use_default_image_token_count, ) @@ -940,7 +810,11 @@ def _count_file_tokens( name_tokens: Final = count_function(filename) if isinstance(filename, str) and filename else 0 if not isinstance(file_data, str) or not file_data: return name_tokens - return name_tokens + _count_opaque_document_tokens(file_data, count_function, use_default_image_token_count) + return name_tokens + calculate_img_tokens( + data=file_data, + mode="auto", + use_default_image_token_count=use_default_image_token_count, + ) def _count_anthropic_content( diff --git a/pyproject.toml b/pyproject.toml index 89c98f8b7f2..2816de2316c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -181,7 +181,7 @@ proxy-runtime = [ "mangum>=0.17.0,<1.0", "azure-ai-contentsafety>=1.0.0,<2.0", "azure-storage-file-datalake>=12.20.0,<13.0", - "pypdf>=6.19.0,<7.0", + "pypdf>=6.16.1,<7.0", "llm-sandbox>=0.3.39,<1.0", "detect-secrets>=1.5.0,<2.0", ] diff --git a/tests/integration/_support/pdf_document.py b/tests/integration/_support/pdf_document.py deleted file mode 100644 index 5c6ad82b85f..00000000000 --- a/tests/integration/_support/pdf_document.py +++ /dev/null @@ -1,149 +0,0 @@ -import base64 -import json -from collections.abc import Iterator, Mapping, Sequence -from dataclasses import dataclass -from itertools import accumulate -from types import MappingProxyType -from typing import Final - -from integration._support.wire import Reply -from pydantic import JsonValue - -PDF_MEDIA_TYPE: Final = "application/pdf" -COUNT_TOKENS_TARGET: Final = "/v1/messages/count_tokens" -COUNT_REFUSED: Final = Reply( - status=400, - body=json.dumps( - {"type": "error", "error": {"type": "invalid_request_error", "message": "count_tokens is not supported"}} - ).encode(), -) - - -@dataclass(frozen=True, slots=True) -class Page: - width: int = 612 - height: int = 792 - text: str | None = None - - -LETTER: Final = Page() -NARROW: Final = Page(width=100, height=1000) -_RENDERED_TOKENS: Final = MappingProxyType({(612, 792): 1534, (100, 1000): 328}) - - -def rendered_tokens(pages: Sequence[Page]) -> int: - return sum(_RENDERED_TOKENS[(page.width, page.height)] for page in pages) - - -def _escaped(text: str) -> str: - return text.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)") - - -def _content_stream(page: Page) -> bytes: - stream: Final = f"BT /F1 12 Tf 72 {page.height - 72} Td ({_escaped(page.text or '')}) Tj ET".encode() - return b"<< /Length " + str(len(stream)).encode() + b" >>\nstream\n" + stream + b"\nendstream" - - -def _page_object(page: Page, contents_id: int | None) -> bytes: - contents: Final = f" /Contents {contents_id} 0 R" if contents_id is not None else "" - return ( - f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {page.width} {page.height}]" - f" /Resources << /Font << /F1 1 0 R >> >>{contents} >>" - ).encode() - - -def _page_bodies(pages: Sequence[Page], starts: Sequence[int]) -> Iterator[bytes]: - for page, start in zip(pages, starts): - if page.text is None: - yield _page_object(page, None) - else: - yield _content_stream(page) - yield _page_object(page, start) - - -def pdf_bytes(pages: Sequence[Page]) -> bytes: - starts: Final = tuple(accumulate((1 if page.text is None else 2 for page in pages), initial=3)) - page_ids: Final = tuple(start + (0 if page.text is None else 1) for start, page in zip(starts, pages)) - kids: Final = " ".join(f"{identity} 0 R" for identity in page_ids) - bodies: Final = ( - b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>", - f"<< /Type /Pages /Kids [{kids}] /Count {len(pages)} >>".encode(), - *_page_bodies(pages, starts), - b"<< /Type /Catalog /Pages 2 0 R >>", - ) - header: Final = b"%PDF-1.4\n" - objects: Final = tuple( - f"{number} 0 obj\n".encode() + body + b"\nendobj\n" for number, body in enumerate(bodies, start=1) - ) - offsets: Final = tuple(accumulate((len(chunk) for chunk in objects), initial=len(header))) - entries: Final = b"".join(f"{offset:010d} 00000 n \n".encode() for offset in offsets[:-1]) - trailer: Final = ( - f"xref\n0 {len(bodies) + 1}\n".encode() - + b"0000000000 65535 f \n" - + entries - + f"trailer\n<< /Size {len(bodies) + 1} /Root {starts[-1]} 0 R >>\nstartxref\n{offsets[-1]}\n%%EOF\n".encode() - ) - return header + b"".join(objects) + trailer - - -def encoded(raw: bytes) -> str: - return base64.b64encode(raw).decode() - - -def data_url(media_type: str, raw: bytes) -> str: - return f"data:{media_type};base64,{encoded(raw)}" - - -def pdf_data_url(pages: Sequence[Page]) -> str: - return data_url(PDF_MEDIA_TYPE, pdf_bytes(pages)) - - -def base64_source(data: JsonValue, media_type: JsonValue = PDF_MEDIA_TYPE) -> dict[str, JsonValue]: - return {"type": "base64", "media_type": media_type, "data": data} - - -def document(source: Mapping[str, JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: - return {"type": "document", "source": dict(source), **fields} - - -def pdf_document(pages: Sequence[Page], **fields: JsonValue) -> dict[str, JsonValue]: - return document(base64_source(encoded(pdf_bytes(pages))), **fields) - - -def text_document(text: str) -> dict[str, JsonValue]: - return document({"type": "text", "media_type": "text/plain", "data": text}) - - -def chat_file(file_data: JsonValue, filename: JsonValue = "document.pdf") -> dict[str, JsonValue]: - return {"type": "file", "file": {"filename": filename, "file_data": file_data}} - - -def responses_input_file(file_data: JsonValue, filename: JsonValue = "document.pdf") -> dict[str, JsonValue]: - return {"type": "input_file", "filename": filename, "file_data": file_data} - - -def messages_body(model: str, blocks: Sequence[JsonValue], text: str, **fields: JsonValue) -> dict[str, JsonValue]: - return { - "model": model, - "max_tokens": 64, - "messages": [{"role": "user", "content": [*blocks, {"type": "text", "text": text}]}], - **fields, - } - - -def chat_body(model: str, parts: Sequence[JsonValue], text: str, **fields: JsonValue) -> dict[str, JsonValue]: - return { - "model": model, - "max_tokens": 64, - "messages": [{"role": "user", "content": [*parts, {"type": "text", "text": text}]}], - **fields, - } - - -def responses_body(model: str, items: Sequence[JsonValue], text: str, **fields: JsonValue) -> dict[str, JsonValue]: - return { - "model": model, - "max_output_tokens": 64, - "input": [{"role": "user", "content": [*items, {"type": "input_text", "text": text}]}], - **fields, - } diff --git a/tests/integration/providers/test_pdf_document_count_tokens_chaos.py b/tests/integration/providers/test_pdf_document_count_tokens_chaos.py deleted file mode 100644 index f196018f117..00000000000 --- a/tests/integration/providers/test_pdf_document_count_tokens_chaos.py +++ /dev/null @@ -1,226 +0,0 @@ -import asyncio -import os -import re -import signal -import threading -import uuid -from collections.abc import Callable -from dataclasses import dataclass -from pathlib import Path -from types import MappingProxyType -from typing import Final -from urllib.parse import urlsplit - -import httpx -import psutil -import pytest -from integration._support.client import Gateway, eventually -from integration._support.database import read_rows -from integration._support.pdf_document import ( - COUNT_REFUSED, - COUNT_TOKENS_TARGET, - LETTER, - messages_body, - pdf_document, - rendered_tokens, -) -from integration._support.process import graceful_stop_seconds, owned_proxy_process -from integration._support.wire import Reply, Request, wire_server -from integration.providers._cache_control_marks_support import anthropic_peer, owned_config -from pydantic import JsonValue, TypeAdapter - -_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") -_MODEL: Final = "burst-claude" -_PROVIDER_KEY: Final = "synthetic-provider-key" -_ASK: Final = "Summarize the attached report in one sentence." -_COUNT_BURST: Final = 24 -_MESSAGE_BURST: Final = 12 -_PAGE_COUNTS: Final = (1, 3, 12) -_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) - - -@dataclass(frozen=True, slots=True) -class _Counted: - pages: int - status: int - text: str - input_tokens: int | None - - -@dataclass(frozen=True, slots=True) -class _Sent: - marker: str - status: int - text: str - call_id: str - - -def _peer(request: Request) -> Reply: - if request.target == COUNT_TOKENS_TARGET: - return COUNT_REFUSED - return anthropic_peer(request) - - -def _held(release: threading.Event) -> Callable[[Request], Reply]: - def respond(request: Request) -> Reply: - assert release.wait(timeout=120), "Held peer was never released" - return _peer(request) - - return respond - - -def _deployment(api_base: str) -> dict[str, JsonValue]: - return { - "model_name": _MODEL, - "litellm_params": {"model": "anthropic/claude-opus-5-5", "api_base": api_base, "api_key": _PROVIDER_KEY}, - } - - -def _count_body(pages: int) -> dict[str, JsonValue]: - blocks: Final[list[JsonValue]] = [pdf_document((LETTER,) * pages)] if pages else [] - return {"model": _MODEL, "messages": messages_body(_MODEL, blocks, _ASK)["messages"]} - - -def _input_tokens(response: httpx.Response) -> int | None: - if response.status_code != 200: - return None - counted: Final = _JSON_OBJECT.validate_json(response.content).get("input_tokens") - return counted if isinstance(counted, int) else None - - -async def _fire_counts(url: str, key: str, *, tolerate_transport_errors: bool = False) -> tuple[_Counted, ...]: - async def one(client: httpx.AsyncClient, index: int) -> _Counted: - pages: Final = _PAGE_COUNTS[index % len(_PAGE_COUNTS)] - response: Final = await client.post( - "/v1/messages/count_tokens", json=_count_body(pages), headers={"Authorization": f"Bearer {key}"} - ) - return _Counted(pages, response.status_code, response.text, _input_tokens(response)) - - async with httpx.AsyncClient(base_url=url, timeout=90, trust_env=False) as client: - results: Final = await asyncio.gather( - *(one(client, index) for index in range(_COUNT_BURST)), return_exceptions=tolerate_transport_errors - ) - for result in results: - assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) - return tuple(result for result in results if isinstance(result, _Counted)) - - -async def _fire_messages(url: str, key: str) -> tuple[_Sent, ...]: - async def one(client: httpx.AsyncClient) -> _Sent: - marker: Final = uuid.uuid4().hex - body: Final = messages_body(_MODEL, [pdf_document((LETTER,))], f"{_ASK} marker-{marker}") - response: Final = await client.post("/v1/messages", json=body, headers={"Authorization": f"Bearer {key}"}) - return _Sent(marker, response.status_code, response.text, response.headers.get("x-litellm-call-id", "")) - - async with httpx.AsyncClient(base_url=url, timeout=90, trust_env=False) as client: - return tuple(await asyncio.gather(*(one(client) for _ in range(_MESSAGE_BURST)))) - - -def _baseline(gateway: Gateway) -> int: - response: Final = gateway.request("POST", "/v1/messages/count_tokens", _count_body(0)) - assert response.status_code == 200, response.text - counted: Final = _input_tokens(response) - assert counted is not None, response.text - return counted - - -def _assert_exact(counts: tuple[_Counted, ...], baseline: int) -> None: - for item in counts: - assert item.status == 200, (item.pages, item.status, item.text) - assert item.input_tokens is not None and item.input_tokens - baseline == rendered_tokens( - (LETTER,) * item.pages - ), ( - item.pages, - item.input_tokens, - baseline, - ) - - -def _single_spend_row(item: _Sent) -> None: - assert item.status == 200, (item.status, item.text) - assert item.call_id, item.text - rows: Final = eventually( - lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (item.call_id,)), - lambda values: len(values) == 1, - seconds=70, - ) - assert len(rows) == 1, item.call_id - - -def _held_upstream_connections(pid: int, upstream: str) -> int: - port: Final = urlsplit(upstream).port - return sum( - 1 - for connection in psutil.Process(pid).net_connections(kind="tcp") - if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port - ) - - -@pytest.mark.timeout(2 * graceful_stop_seconds() + 180) -async def test_pdf_count_burst_across_two_workers_prices_every_page_and_logs_each_message_once( - gateway: Gateway, tmp_path: Path -) -> None: - with wire_server(_peer) as wire: - config: Final = owned_config(tmp_path, [_deployment(wire.url)], litellm_settings={"cache": False}) - with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: - eventually( - lambda: _STARTED_WORKER.findall(owned.log.read_text()), - lambda pids: len(pids) == 2, - seconds=graceful_stop_seconds(), - ) - owned_url: Final = str(owned.gateway.client.base_url) - baseline: Final = _baseline(owned.gateway) - wire.drain() - counts, messages = await asyncio.gather( - _fire_counts(owned_url, owned.gateway.key), _fire_messages(owned_url, owned.gateway.key) - ) - received: Final = wire.drain() - targets: Final = [request.target for request in received] - assert targets.count(COUNT_TOKENS_TARGET) == _COUNT_BURST, targets - assert targets.count("/v1/messages") == _MESSAGE_BURST, targets - _assert_exact(counts, baseline) - assert {item.marker for item in messages} == { - _JSON_OBJECT.validate_json(item.text)["id"][len("msg_") :] for item in messages if item.status == 200 - } - for item in messages: - _single_spend_row(item) - - -@pytest.mark.timeout(2 * graceful_stop_seconds() + 180) -async def test_worker_sigkill_mid_pdf_count_burst_leaves_the_sibling_pricing_pages( - gateway: Gateway, tmp_path: Path -) -> None: - release: Final = threading.Event() - with wire_server(_held(release)) as wire: - config: Final = owned_config(tmp_path, [_deployment(wire.url)], litellm_settings={"cache": False}) - with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: - workers: Final[tuple[int, ...]] = eventually( - lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), - lambda pids: len(pids) == 2, - seconds=graceful_stop_seconds(), - ) - owned_url: Final = str(owned.gateway.client.base_url) - burst: Final = asyncio.create_task( - _fire_counts(owned_url, owned.gateway.key, tolerate_transport_errors=True) - ) - try: - await asyncio.to_thread( - eventually, lambda: wire.received.qsize(), lambda size: size >= _COUNT_BURST, 90 - ) - held_by: Final = MappingProxyType({pid: _held_upstream_connections(pid, wire.url) for pid in workers}) - victim: Final = max(workers, key=held_by.__getitem__) - os.kill(victim, signal.SIGKILL) - finally: - release.set() - served: Final = await burst - wire.drain() - baseline: Final = _baseline(owned.gateway) - after: Final = await _fire_counts(owned_url, owned.gateway.key) - after_received: Final = wire.drain() - assert sum(held_by.values()) == _COUNT_BURST, held_by - assert held_by[victim] > 0, held_by - assert len(served) == _COUNT_BURST - held_by[victim], (len(served), held_by) - assert len(after) == _COUNT_BURST, len(after) - assert [request.target for request in after_received] == [COUNT_TOKENS_TARGET] * (_COUNT_BURST + 1) - _assert_exact(served, baseline) - _assert_exact(after, baseline) diff --git a/tests/integration/providers/test_pdf_document_count_tokens_wire.py b/tests/integration/providers/test_pdf_document_count_tokens_wire.py deleted file mode 100644 index 41d5f3d8685..00000000000 --- a/tests/integration/providers/test_pdf_document_count_tokens_wire.py +++ /dev/null @@ -1,273 +0,0 @@ -import io -import json -import uuid -from collections.abc import Sequence -from types import MappingProxyType -from typing import Final - -import pytest -from integration._support.client import Gateway, eventually, object_value -from integration._support.database import read_rows -from integration._support.pdf_document import ( - COUNT_REFUSED, - COUNT_TOKENS_TARGET, - LETTER, - NARROW, - Page, - base64_source, - chat_file, - document, - encoded, - messages_body, - pdf_bytes, - pdf_data_url, - pdf_document, - rendered_tokens, - responses_input_file, - text_document, -) -from integration._support.provider import SharedProvider -from integration._support.wire import Reply -from pydantic import JsonValue, TypeAdapter -from pypdf import PdfReader, PdfWriter - -_MODEL: Final = "anthropic/claude-opus-5-5" -_ASK: Final = "Summarize the attached report in one sentence." -_TEXT: Final = "The quarterly report covers revenue, margins and headcount." -_TITLE: Final = "Quarterly report" -_CONTEXT: Final = "Board pack, page one" -_PEER_INPUT_TOKENS: Final = 12 -_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) -_PNG_HEADER: Final = b"\x89PNG\r\n\x1a\n\x00\x00\x00\x0dIHDR" + (1568).to_bytes(4, "big") * 2 + b"\x08\x06\x00\x00\x00" - - -def _locked(raw: bytes, user_password: str) -> bytes: - writer: Final = PdfWriter(clone_from=PdfReader(io.BytesIO(raw))) - writer.encrypt(user_password=user_password, owner_password="owner", algorithm="RC4-128") - out: Final = io.BytesIO() - writer.write(out) - return out.getvalue() - - -_UNREADABLE: Final[MappingProxyType[str, JsonValue]] = MappingProxyType( - { - "garbage-5kb": encoded(bytes(range(256)) * 20), - "cut-mid-stream": encoded(pdf_bytes((LETTER, LETTER))[:200]), - "png-bytes": encoded(_PNG_HEADER), - "user-password": encoded(_locked(pdf_bytes((LETTER, LETTER)), "reader")), - "empty": "", - "int": 1234, - } -) - - -def _int(value: JsonValue) -> int: - assert isinstance(value, int), value - return value - - -def _turn(blocks: Sequence[JsonValue]) -> list[JsonValue]: - return messages_body(_MODEL, blocks, _ASK)["messages"] - - -def _count(gateway: Gateway, provider: SharedProvider, blocks: Sequence[JsonValue]) -> int: - provider.expect(COUNT_REFUSED) - response: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": _MODEL, "messages": _turn(blocks)}) - assert response.status_code == 200, response.text - assert [(request.method, request.target) for request in provider.received()] == [("POST", COUNT_TOKENS_TARGET)] - return _int(_JSON_OBJECT.validate_json(response.content)["input_tokens"]) - - -def _local(gateway: Gateway, parts: Sequence[JsonValue]) -> int: - response: Final = gateway.request( - "POST", "/utils/token_counter", {"model": _MODEL, "messages": _turn(parts)}, params={"call_endpoint": "false"} - ) - assert response.status_code == 200, response.text - return _int(_JSON_OBJECT.validate_json(response.content)["total_tokens"]) - - -def _responses_count(gateway: Gateway, provider: SharedProvider, items: Sequence[JsonValue]) -> int: - provider.expect(COUNT_REFUSED) - response: Final = gateway.request( - "POST", - "/v1/responses/input_tokens", - {"model": _MODEL, "input": [{"role": "user", "content": [*items, {"type": "input_text", "text": _ASK}]}]}, - ) - assert response.status_code == 200, response.text - assert [(request.method, request.target) for request in provider.received()] == [("POST", COUNT_TOKENS_TARGET)] - payload: Final = _JSON_OBJECT.validate_json(response.content) - assert payload["object"] == "response.input_tokens", response.text - return _int(payload["input_tokens"]) - - -def _cost(gateway: Gateway, parts: Sequence[JsonValue]) -> float: - payload: Final = gateway.post("/spend/calculate", {"model": _MODEL, "messages": _turn(parts)}) - cost: Final = payload["cost"] - assert isinstance(cost, (int, float)), payload - return float(cost) - - -def _input_price(gateway: Gateway) -> float: - listed: Final = gateway.get("/model/info")["data"] - assert isinstance(listed, list), listed - rows: Final = [object_value(row) for row in listed if isinstance(row, dict) and row.get("model_name") == _MODEL] - assert len(rows) == 1, rows - price: Final = object_value(rows[0]["model_info"])["input_cost_per_token"] - assert isinstance(price, float) and price > 0, price - return price - - -def _peer_message(identity: str) -> Reply: - return Reply( - body=json.dumps( - { - "id": identity, - "type": "message", - "role": "assistant", - "model": "claude-opus-5-5", - "content": [{"type": "text", "text": "One sentence."}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": _PEER_INPUT_TOKENS, "output_tokens": 3}, - } - ).encode() - ) - - -def _prompt_tokens(call_id: str) -> int: - rows: Final = eventually( - lambda: read_rows('SELECT prompt_tokens FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (call_id,)), - lambda values: len(values) == 1, - seconds=70, - ) - return _int(rows[0]["prompt_tokens"]) - - -@pytest.mark.parametrize("pages", [1, 3, 12]) -def test_count_tokens_fallback_prices_each_blank_page_as_a_rendered_image( - gateway: Gateway, provider: SharedProvider, pages: int -) -> None: - letter_pages: Final = (LETTER,) * pages - with_document: Final = _count(gateway, provider, [pdf_document(letter_pages)]) - without: Final = _count(gateway, provider, []) - assert with_document - without == rendered_tokens(letter_pages), (with_document, without) - - -def test_count_tokens_fallback_adds_a_page_text_to_its_rendered_image( - gateway: Gateway, provider: SharedProvider -) -> None: - text_page: Final = _count(gateway, provider, [pdf_document((Page(text=_TEXT),))]) - blank_page: Final = _count(gateway, provider, [pdf_document((LETTER,))]) - text_source: Final = _count(gateway, provider, [text_document(_TEXT)]) - without: Final = _count(gateway, provider, []) - assert text_source - without > 0, (text_source, without) - assert text_page - blank_page == text_source - without, (text_page, blank_page, text_source, without) - - -def test_count_tokens_fallback_prices_a_page_by_its_rendered_area(gateway: Gateway, provider: SharedProvider) -> None: - mixed: Final = (LETTER, NARROW, LETTER) - without: Final = _count(gateway, provider, []) - narrow: Final = _count(gateway, provider, [pdf_document((NARROW,))]) - assert narrow - without == rendered_tokens((NARROW,)), (narrow, without) - assert _count(gateway, provider, [pdf_document(mixed)]) - without == rendered_tokens(mixed) - - -def test_count_tokens_fallback_prices_every_document_in_the_turn(gateway: Gateway, provider: SharedProvider) -> None: - both: Final = _count(gateway, provider, [pdf_document((LETTER,) * 3), pdf_document((LETTER,))]) - without: Final = _count(gateway, provider, []) - assert both - without == rendered_tokens((LETTER,) * 4), (both, without) - - -def test_count_tokens_fallback_prices_a_pdf_without_pages_as_nothing( - gateway: Gateway, provider: SharedProvider -) -> None: - assert _count(gateway, provider, [pdf_document(())]) == _count(gateway, provider, []) - - -def test_count_tokens_fallback_adds_title_and_context_once_per_document( - gateway: Gateway, provider: SharedProvider -) -> None: - without: Final = _count(gateway, provider, []) - plain: Final = _count(gateway, provider, [pdf_document((LETTER,))]) - annotated: Final = _count(gateway, provider, [pdf_document((LETTER,), title=_TITLE, context=_CONTEXT)]) - title: Final = _count(gateway, provider, [text_document(_TITLE)]) - without - context: Final = _count(gateway, provider, [text_document(_CONTEXT)]) - without - assert title > 0 and context > 0, (title, context) - assert annotated - plain == title + context, (annotated, plain, title, context) - - -@pytest.mark.parametrize("label", list(_UNREADABLE)) -def test_count_tokens_fallback_prices_unreadable_pdf_bytes_like_one_image( - gateway: Gateway, provider: SharedProvider, label: str -) -> None: - data: Final = _UNREADABLE[label] - as_pdf: Final = _count(gateway, provider, [document(base64_source(data))]) - as_png: Final = _count(gateway, provider, [document(base64_source(data, media_type="image/png"))]) - without: Final = _count(gateway, provider, []) - assert as_pdf == as_png, (as_pdf, as_png) - assert 0 <= as_pdf - without < rendered_tokens((LETTER,)), (as_pdf, without) - - -def test_count_tokens_fallback_reads_a_list_wrapped_base64_string_like_the_bare_string( - gateway: Gateway, provider: SharedProvider -) -> None: - raw: Final = encoded(pdf_bytes((LETTER,))) - wrapped: Final = _count(gateway, provider, [document(base64_source([raw]))]) - assert wrapped == _count(gateway, provider, [document(base64_source(raw))]), wrapped - - -def test_count_tokens_fallback_reads_an_owner_locked_pdf(gateway: Gateway, provider: SharedProvider) -> None: - pages: Final = (LETTER, LETTER) - locked: Final = document(base64_source(encoded(_locked(pdf_bytes(pages), "")))) - assert _count(gateway, provider, [locked]) - _count(gateway, provider, []) == rendered_tokens(pages) - - -def test_count_tokens_fallback_parses_only_the_pdf_media_type(gateway: Gateway, provider: SharedProvider) -> None: - raw: Final = encoded(pdf_bytes((LETTER,))) - without: Final = _count(gateway, provider, []) - labelled: Final = _count(gateway, provider, [document(base64_source(raw))]) - assert labelled - without == rendered_tokens((LETTER,)), (labelled, without) - mislabelled: Final = _count(gateway, provider, [document(base64_source(raw, media_type="application/x-pdf"))]) - assert mislabelled == _count(gateway, provider, [document(base64_source(raw, media_type="image/png"))]) - - -def test_utils_token_counter_prices_a_chat_file_by_its_pages(gateway: Gateway) -> None: - one: Final = _local(gateway, [chat_file(pdf_data_url((LETTER,)))]) - three: Final = _local(gateway, [chat_file(pdf_data_url((LETTER,) * 3))]) - assert three - one == rendered_tokens((LETTER,) * 2), (three, one) - text_page: Final = _local(gateway, [chat_file(pdf_data_url((Page(text=_TEXT),)))]) - text_part: Final = _local(gateway, [chat_file(pdf_data_url((LETTER,))), {"type": "text", "text": _TEXT}]) - assert text_part - one > 0, (text_part, one) - assert text_page - one == text_part - one, (text_page, text_part, one) - - -def test_responses_input_tokens_fallback_prices_an_input_file_by_its_pages( - gateway: Gateway, provider: SharedProvider -) -> None: - one: Final = _responses_count(gateway, provider, [responses_input_file(pdf_data_url((LETTER,)))]) - three: Final = _responses_count(gateway, provider, [responses_input_file(pdf_data_url((LETTER,) * 3))]) - assert three - one == rendered_tokens((LETTER,) * 2), (three, one) - - -def test_spend_calculate_prices_a_chat_file_by_its_pages(gateway: Gateway) -> None: - twelve: Final = _cost(gateway, [chat_file(pdf_data_url((LETTER,) * 12))]) - one: Final = _cost(gateway, [chat_file(pdf_data_url((LETTER,)))]) - assert twelve - one == pytest.approx(rendered_tokens((LETTER,) * 11) * _input_price(gateway)), (twelve, one) - - -def test_a_cached_pdf_message_keeps_the_peer_usage_on_both_spend_rows( - gateway: Gateway, provider: SharedProvider -) -> None: - marker: Final = uuid.uuid4().hex - body: Final = messages_body(_MODEL, [pdf_document((LETTER,))], f"{_ASK} marker-{marker}") - provider.expect(_peer_message(f"msg_{marker}")) - first: Final = gateway.request("POST", "/v1/messages", body) - second: Final = gateway.request("POST", "/v1/messages", body) - assert (first.status_code, second.status_code) == (200, 200), (first.text, second.text) - assert [(request.method, request.target) for request in provider.received()] == [("POST", "/v1/messages")] - assert _JSON_OBJECT.validate_json(second.content)["id"] == f"msg_{marker}", second.text - assert first.headers["x-litellm-call-id"] != second.headers["x-litellm-call-id"] - assert [_prompt_tokens(response.headers["x-litellm-call-id"]) for response in (first, second)] == [ - _PEER_INPUT_TOKENS, - _PEER_INPUT_TOKENS, - ] diff --git a/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py b/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py deleted file mode 100644 index 9fad73f8080..00000000000 --- a/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py +++ /dev/null @@ -1,292 +0,0 @@ -import os -import socket -import threading -import uuid -from collections.abc import Callable, Iterator, Mapping -from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass -from pathlib import Path -from typing import Final - -import httpx -import pytest -from integration._support.client import Gateway, eventually, gateway_from_environment, object_value -from integration._support.database import read_rows, scratch_database -from integration._support.pdf_document import ( - COUNT_REFUSED, - COUNT_TOKENS_TARGET, - LETTER, - chat_body, - chat_file, - messages_body, - pdf_data_url, - pdf_document, - responses_body, - responses_input_file, -) -from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process -from integration._support.wire import Reply, Request, Wire, wire_server -from integration.providers._cache_control_marks_support import anthropic_peer, marker_of, owned_config -from pydantic import JsonValue, TypeAdapter -from redis import Redis - -_WINDOW: Final = "window-claude" -_ITPM: Final = "itpm-claude" -_BUDGET: Final = "budget-claude" -_AFFINITY: Final = "affinity-claude" -_LIMIT: Final = 5000 -_PAGE_COUNT: Final = 12 -_PAGES: Final = (LETTER,) * _PAGE_COUNT -_ASK: Final = "Summarize the attached report in one sentence." -_FOLLOW_UPS: Final = 9 -_TRANSCRIPT: Final = " ".join(f"Line {number}: revenue, margins and headcount moved." for number in range(160)) -_PIN_KEYS: Final = "*:prompt_caching" -_PROVIDER_KEY: Final = "synthetic-provider-key" -_SURFACES: Final = ("messages", "messages-stream", "chat-file", "responses-input-file") -_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) - - -@dataclass(frozen=True, slots=True) -class _Rig: - owned: OwnedProxy - wire: Wire - held_port: int - - -def _marker() -> str: - return uuid.uuid4().hex - - -def _free_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return int(reserve.getsockname()[1]) - - -def _deployment( - name: str, - backend: str, - api_base: str, - *, - model_info: Mapping[str, JsonValue] | None = None, - **params: JsonValue, -) -> dict[str, JsonValue]: - return { - "model_name": name, - "litellm_params": {"model": f"anthropic/{backend}", "api_base": api_base, "api_key": _PROVIDER_KEY, **params}, - "model_info": dict(model_info or {}), - } - - -def _peer(request: Request) -> Reply: - if request.target == COUNT_TOKENS_TARGET: - return COUNT_REFUSED - return anthropic_peer(request) - - -def _held(release: threading.Event) -> Callable[[Request], Reply]: - def respond(request: Request) -> Reply: - assert release.wait(timeout=120), "Held peer was never released" - return _peer(request) - - return respond - - -def _request(surface: str, model: str, pages: int, marker: str) -> tuple[str, dict[str, JsonValue]]: - text: Final = f"{_ASK} marker-{marker}" - letter_pages: Final = (LETTER,) * pages - if surface == "messages": - return "/v1/messages", messages_body(model, [pdf_document(letter_pages)], text) - if surface == "messages-stream": - return "/v1/messages", messages_body(model, [pdf_document(letter_pages)], text, stream=True) - if surface == "chat-file": - return "/v1/chat/completions", chat_body(model, [chat_file(pdf_data_url(letter_pages))], text) - return "/v1/responses", responses_body(model, [responses_input_file(pdf_data_url(letter_pages))], text) - - -@pytest.fixture(scope="module") -def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: - directory: Final = tmp_path_factory.mktemp("pdf-pre-call") - held_port: Final = _free_port() - with gateway_from_environment() as gateway, wire_server(_peer) as wire: - config: Final = owned_config( - directory, - [ - _deployment(_WINDOW, "claude-haiku-5-5", wire.url, model_info={"max_input_tokens": _LIMIT}), - _deployment(_ITPM, "claude-sonnet-4-6", wire.url, itpm=_LIMIT), - _deployment(_BUDGET, "claude-opus-4-8", f"http://127.0.0.1:{held_port}"), - ], - litellm_settings={"cache": False}, - router_settings={ - "enable_pre_call_checks": True, - "optional_pre_call_checks": ["enforce_model_rate_limits"], - }, - ) - with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned: - yield _Rig(owned, wire, held_port) - - -def _spend_row(call_id: str) -> dict[str, JsonValue]: - rows: Final = eventually( - lambda: read_rows('SELECT status FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (call_id,)), - lambda values: len(values) == 1, - seconds=70, - ) - return rows[0] - - -@pytest.mark.parametrize("surface", _SURFACES) -def test_context_window_check_rejects_a_pdf_whose_pages_exceed_max_input_tokens(rig: _Rig, surface: str) -> None: - rig.wire.drain() - path, body = _request(surface, _WINDOW, _PAGE_COUNT, _marker()) - response: Final = rig.owned.gateway.request("POST", path, body) - assert response.status_code == 400, (response.status_code, response.text) - assert "Context Window exceeded" in response.text, response.text - assert [request.target for request in rig.wire.drain()] == [] - - -def test_context_window_check_passes_a_one_page_pdf_and_forwards_its_bytes(rig: _Rig) -> None: - rig.wire.drain() - path, body = _request("messages", _WINDOW, 1, _marker()) - response: Final = rig.owned.gateway.request("POST", path, body) - assert response.status_code == 200, response.text - received: Final = rig.wire.drain() - assert [request.target for request in received] == ["/v1/messages"] - forwarded: Final = _JSON_OBJECT.validate_json(received[0].body)["messages"] - assert isinstance(forwarded, list) and len(forwarded) == 1, forwarded - blocks: Final = object_value(forwarded[0])["content"] - assert isinstance(blocks, list) and pdf_document((LETTER,)) in blocks, blocks - - -def test_model_itpm_check_rejects_a_pdf_whose_pages_exceed_the_limit(rig: _Rig) -> None: - rig.wire.drain() - path, body = _request("messages", _ITPM, _PAGE_COUNT, _marker()) - response: Final = rig.owned.gateway.request("POST", path, body) - assert response.status_code == 429, (response.status_code, response.text) - assert '"error"' in response.text, response.text - assert [request.target for request in rig.wire.drain()] == [] - - -def test_key_budget_reservation_rejects_the_second_concurrent_pdf_request(rig: _Rig) -> None: - release: Final = threading.Event() - requests: Final = tuple(_request("messages", _BUDGET, _PAGE_COUNT, _marker()) for _ in range(2)) - with wire_server(_held(release), port=rig.held_port) as wire, rig.owned.gateway.scenario() as scenario: - key: Final = scenario.key(max_budget=0.01) - headers: Final = {"Authorization": f"Bearer {key}"} - with ( - httpx.Client(base_url=str(rig.owned.gateway.client.base_url), timeout=90, trust_env=False) as client, - ThreadPoolExecutor(max_workers=2) as pool, - ): - futures: Final = tuple( - pool.submit(client.post, path, json=body, headers=headers) for path, body in requests - ) - try: - eventually( - lambda: wire.received.qsize() + sum(1 for future in futures if future.done()), - lambda settled: settled >= 2, - seconds=60, - ) - finally: - release.set() - responses: Final = tuple(future.result(timeout=90) for future in futures) - received: Final = wire.drain() - assert sorted(response.status_code for response in responses) == [200, 422], [ - response.text for response in responses - ] - rejected: Final = next(response for response in responses if response.status_code != 200) - assert "udget" in rejected.text, rejected.text - assert [request.target for request in received] == ["/v1/messages"] - - -def test_a_pre_call_rejection_logs_one_failure_row_and_calls_no_upstream(rig: _Rig) -> None: - path, body = _request("messages", _WINDOW, _PAGE_COUNT, _marker()) - response: Final = rig.owned.gateway.request("POST", path, body) - assert response.status_code == 400, (response.status_code, response.text) - assert _spend_row(response.headers["x-litellm-call-id"])["status"] == "failure" - assert rig.wire.drain() == () - - -def _first_turn(session: str, marker: str) -> list[JsonValue]: - return [ - {"role": "system", "content": f"Answer from the attached report. session {session}"}, - *chat_body( - _AFFINITY, - [{"type": "text", "text": _TRANSCRIPT}, chat_file(pdf_data_url(_PAGES))], - f"{_ASK} marker-{marker}", - )["messages"], - ] - - -def _follow_up(first_turn: list[JsonValue], marker: str) -> list[JsonValue]: - return [ - *first_turn, - {"role": "assistant", "content": "One sentence."}, - {"role": "user", "content": f"And the margins? marker-{marker}"}, - ] - - -def _served(wire: Wire) -> frozenset[str]: - return frozenset(marker_of(request) for request in wire.drain()) - - -def _pin_count(cache: Redis) -> int: - return len(cache.keys(_PIN_KEYS)) - - -@pytest.mark.timeout(2 * graceful_stop_seconds() + 180) -def test_prompt_caching_affinity_pins_follow_ups_of_a_cached_turn_that_carries_a_pdf( - gateway: Gateway, tmp_path: Path -) -> None: - session: Final = _marker() - first_marker: Final = _marker() - follow_up_markers: Final = tuple(_marker() for _ in range(_FOLLOW_UPS)) - first_turn: Final = _first_turn(session, first_marker) - with scratch_database() as database_url, wire_server(_peer) as left, wire_server(_peer) as right: - config: Final = owned_config( - tmp_path, - [ - _deployment(_AFFINITY, "claude-opus-5-5", left.url, model_info={"id": f"pdf-left-{session}"}), - _deployment(_AFFINITY, "claude-opus-5-5", right.url, model_info={"id": f"pdf-right-{session}"}), - ], - litellm_settings={"cache": False}, - router_settings={ - "optional_pre_call_checks": ["prompt_caching"], - "redis_host": os.environ["REDIS_HOST"], - "redis_port": int(os.environ["REDIS_PORT"]), - }, - ) - with ( - Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache, - owned_proxy_process( - gateway, - tmp_path, - {"DATABASE_URL": database_url}, - config=config, - remove_environment=("DATABASE_URL_READ_REPLICA",), - workers=2, - ) as owned, - owned.gateway.scenario() as scenario, - ): - key: Final = scenario.key(metadata={"enable_prompt_caching": True}) - pins_before: Final = _pin_count(cache) - first: Final = owned.gateway.request( - "POST", "/v1/chat/completions", {"model": _AFFINITY, "messages": first_turn, "max_tokens": 64}, key=key - ) - assert first.status_code == 200, first.text - eventually(lambda: _pin_count(cache), lambda count: count > pins_before, seconds=70) - follow_ups: Final = tuple( - owned.gateway.request( - "POST", - "/v1/chat/completions", - {"model": _AFFINITY, "messages": _follow_up(first_turn, marker), "max_tokens": 64}, - key=key, - ) - for marker in follow_up_markers - ) - served: Final = {"left": _served(left), "right": _served(right)} - assert [response.status_code for response in follow_ups] == [200] * _FOLLOW_UPS, [ - response.text for response in follow_ups - ] - assert served["left"] | served["right"] == {first_marker, *follow_up_markers}, served - first_side: Final = next(side for side in served if first_marker in served[side]) - assert served[first_side] == {first_marker, *follow_up_markers}, served diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 4b260c387b7..8d53e14388b 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -3,14 +3,12 @@ import asyncio import base64 import importlib -import io import json import os import subprocess import sys import threading import time -import zlib from collections.abc import Mapping from concurrent.futures import Future, wait from itertools import accumulate, chain @@ -1365,17 +1363,54 @@ def _count_user_content(content: list[dict]) -> int: ) +def _pdf_base64(pages: tuple[str, ...], width: int = 612, height: int = 792) -> str: + def page_objects(index: int, text: str) -> tuple[bytes, bytes]: + escaped: Final = text.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)") + stream: Final = f"BT /F1 12 Tf 72 720 Td ({escaped}) Tj ET".encode("latin-1") + content: Final = f"<< /Length {len(stream)} >>\nstream\n".encode() + stream + b"\nendstream" + page: Final = ( + f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {width} {height}] " + f"/Resources << /Font << /F1 3 0 R >> >> /Contents {4 + 2 * index} 0 R >>" + ).encode() + return content, page + + kids: Final = " ".join(f"{5 + 2 * index} 0 R" for index in range(len(pages))) + bodies: Final = ( + b"<< /Type /Catalog /Pages 2 0 R >>", + f"<< /Type /Pages /Count {len(pages)} /Kids [ {kids} ] >>".encode(), + b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>", + *chain.from_iterable(page_objects(index, text) for index, text in enumerate(pages)), + ) + header: Final = b"%PDF-1.4\n" + objects: Final = tuple( + f"{number} 0 obj\n".encode() + body + b"\nendobj\n" for number, body in enumerate(bodies, start=1) + ) + offsets: Final = accumulate((len(header), *(len(obj) for obj in objects[:-1]))) + xref: Final = f"xref\n0 {len(objects) + 1}\n0000000000 65535 f \n".encode() + b"".join( + f"{offset:010d} 00000 n \n".encode() for offset in offsets + ) + trailer: Final = ( + f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\n" + f"startxref\n{len(header) + sum(len(obj) for obj in objects)}\n%%EOF\n" + ).encode() + return base64.b64encode(header + b"".join(objects) + xref + trailer).decode() + + +_READABLE_PDF_BASE64: Final = _pdf_base64(("Count this document without parsing its text locally.",)) + + @pytest.mark.parametrize( "source", [ {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, {"type": "url", "url": "https://example.com/report.pdf"}, {"type": "file", "file_id": "file-abc123"}, + {"type": "base64", "media_type": "application/pdf", "data": _READABLE_PDF_BASE64}, ], - ids=["base64", "url", "file"], + ids=["base64", "url", "file", "readable_pdf"], ) def test_anthropic_document_block_with_opaque_source_is_priced_like_an_image(source: dict[str, str]): - """A `document` with no readable pages (a bare PDF header, a URL, a file id) is priced like an `image`, not raised on.""" + """A `document` whose bytes can't be tokenized locally is priced like an `image`, not raised on.""" prompt = {"type": "text", "text": "Summarize this file."} assert _count_user_content([prompt, {"type": "document", "source": source}]) == _count_user_content( @@ -1414,7 +1449,8 @@ def test_anthropic_document_title_and_context_add_their_tokens(): ) -def test_openai_file_block_prices_like_the_equivalent_anthropic_document(): +@pytest.mark.parametrize("pdf_data", ("JVBERi0xLjQK", _READABLE_PDF_BASE64), ids=("opaque", "readable")) +def test_openai_file_block_prices_like_the_equivalent_anthropic_document(pdf_data: str) -> None: """An inline `file` is a `document` in the chat-completions dialect, so it must price identically, not raise. Before the fix `file` was missing from the content-block match even though `ChatCompletionFileObject` @@ -1424,22 +1460,20 @@ def test_openai_file_block_prices_like_the_equivalent_anthropic_document(): prompt = {"type": "text", "text": "Summarize this file."} inline_file = { "type": "file", - "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"}, + "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64," + pdf_data}, } document = { "type": "document", "title": "report.pdf", - "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + "source": {"type": "base64", "media_type": "application/pdf", "data": pdf_data}, } assert _count_user_content([prompt, inline_file]) == _count_user_content([prompt, document]) assert _count_user_content([prompt, inline_file]) > _count_user_content([prompt]) - readable: Final = _pdf_base64(("Revenue grew eleven percent while churn fell to two percent.",)) - readable_file: Final = {"type": "file", "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64," + readable}} - readable_document: Final = {"type": "document", "title": "report.pdf", "source": _pdf_source(readable)} - assert _count_user_content([prompt, readable_file]) == _count_user_content([prompt, readable_document]) - assert _count_user_content([prompt, readable_file]) > _count_user_content([prompt, inline_file]) + assert _count_user_content([prompt, inline_file]) == _count_user_content( + [prompt, {"type": "text", "text": "report.pdf"}, {"type": "image", "source": document["source"]}] + ) def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): @@ -1453,222 +1487,6 @@ def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): assert _count_user_content([prompt, named]) == _count_user_content([prompt, {"type": "text", "text": "report.pdf"}]) -def _pdf_base64( - pages: tuple[str, ...], - width: int = 612, - height: int = 792, - *, - padding: int = 0, - compressed: bool = False, - unused_objects: int = 0, - unused_bytes: int = 0, -) -> str: - def page_objects(index: int, text: str) -> tuple[bytes, bytes]: - escaped: Final = text.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)") - raw_stream: Final = b" " * padding + f"BT /F1 12 Tf 72 720 Td ({escaped}) Tj ET".encode("latin-1") - stream: Final = zlib.compress(raw_stream) if compressed else raw_stream - filter_field: Final = " /Filter /FlateDecode" if compressed else "" - content: Final = f"<< /Length {len(stream)}{filter_field} >>\nstream\n".encode() + stream + b"\nendstream" - page: Final = ( - f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {width} {height}] " - f"/Resources << /Font << /F1 3 0 R >> >> /Contents {4 + 2 * index} 0 R >>" - ).encode() - return content, page - - kids: Final = " ".join(f"{5 + 2 * index} 0 R" for index in range(len(pages))) - bodies: Final = ( - b"<< /Type /Catalog /Pages 2 0 R >>", - f"<< /Type /Pages /Count {len(pages)} /Kids [ {kids} ] >>".encode(), - b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>", - *chain.from_iterable(page_objects(index, text) for index, text in enumerate(pages)), - *(b"0" for _ in range(unused_objects)), - *((b"(" + b" " * unused_bytes + b")",) if unused_bytes else ()), - ) - header: Final = b"%PDF-1.4\n" - objects: Final = tuple( - f"{number} 0 obj\n".encode() + body + b"\nendobj\n" for number, body in enumerate(bodies, start=1) - ) - offsets: Final = accumulate((len(header), *(len(obj) for obj in objects[:-1]))) - xref: Final = f"xref\n0 {len(objects) + 1}\n0000000000 65535 f \n".encode() + b"".join( - f"{offset:010d} 00000 n \n".encode() for offset in offsets - ) - trailer: Final = ( - f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\n" - f"startxref\n{len(header) + sum(len(obj) for obj in objects)}\n%%EOF\n" - ).encode() - return base64.b64encode(header + b"".join(objects) + xref + trailer).decode() - - -def _pdf_source(pdf_base64: str) -> dict[str, str]: - return {"type": "base64", "media_type": "application/pdf", "data": pdf_base64} - - -@pytest.mark.parametrize("case", ("pages", "compressed_stream", "total_streams", "objects", "input_bytes")) -def test_pdf_counting_over_budget_preserves_legacy_image_estimate(case: str) -> None: - from litellm.litellm_core_utils.token_counter import ( - _PDF_MAX_INPUT_BYTES, - _PDF_MAX_OBJECTS, - _PDF_MAX_PAGES, - _PDF_MAX_STREAM_BYTES, - _PDF_MAX_TOTAL_STREAM_BYTES, - ) - - pdf: Final = { - "pages": lambda: _pdf_base64(("Count this",) * (_PDF_MAX_PAGES + 1)), - "compressed_stream": lambda: _pdf_base64(("Count this",), padding=_PDF_MAX_STREAM_BYTES, compressed=True), - "total_streams": lambda: _pdf_base64( - ("Count this",) * (_PDF_MAX_TOTAL_STREAM_BYTES // _PDF_MAX_STREAM_BYTES + 1), - padding=_PDF_MAX_STREAM_BYTES - 100, - compressed=True, - ), - "objects": lambda: _pdf_base64(("Count this",), unused_objects=_PDF_MAX_OBJECTS), - "input_bytes": lambda: _pdf_base64(("Count this",), unused_bytes=_PDF_MAX_INPUT_BYTES), - }[case]() - source: Final = _pdf_source(pdf) - assert _count_user_content([{"type": "document", "source": source}]) == _count_user_content( - [{"type": "image", "source": source}] - ) - - -def test_compressed_pdf_within_budget_keeps_text_and_page_counts() -> None: - pages: Final = ("Count the first page", "Count the second page") - compressed: Final = _pdf_source(_pdf_base64(pages, compressed=True)) - uncompressed: Final = _pdf_source(_pdf_base64(pages)) - assert _count_user_content([{"type": "document", "source": compressed}]) == _count_user_content( - [{"type": "document", "source": uncompressed}] - ) - - -@pytest.mark.parametrize("nested", (False, True)) -def test_pdf_form_budget_never_returns_a_partial_count(nested: bool) -> None: - from pypdf import PdfWriter - from pypdf.generic import ArrayObject, DecodedStreamObject, DictionaryObject, NameObject, NumberObject - - writer: Final = PdfWriter() - page: Final = writer.add_blank_page(width=612, height=792) - font: Final = DictionaryObject( - { - NameObject("/Type"): NameObject("/Font"), - NameObject("/Subtype"): NameObject("/Type1"), - NameObject("/BaseFont"): NameObject("/Helvetica"), - } - ) - form: Final = DecodedStreamObject() - form.set_data(b"BT /F1 12 Tf 72 720 Td (Every form must be counted) Tj ET") - form.update( - { - NameObject("/Type"): NameObject("/XObject"), - NameObject("/Subtype"): NameObject("/Form"), - NameObject("/BBox"): ArrayObject([NumberObject(0), NumberObject(0), NumberObject(612), NumberObject(792)]), - NameObject("/Resources"): DictionaryObject( - {NameObject("/Font"): DictionaryObject({NameObject("/F1"): font})} - ), - } - ) - resources: Final = DictionaryObject( - {NameObject("/XObject"): DictionaryObject({NameObject("/Fm"): writer._add_object(form)})} - ) - repeated: Final = DecodedStreamObject() - repeated.set_data(b"/Fm Do\n" * 17) - repeated.update( - { - NameObject("/Subtype"): NameObject("/Form"), - NameObject("/BBox"): form["/BBox"], - NameObject("/Resources"): resources, - } - ) - outer: Final = DecodedStreamObject() - outer.set_data(b"/Fm Do") - page[NameObject("/Resources")] = ( - DictionaryObject({NameObject("/XObject"): DictionaryObject({NameObject("/Fm"): writer._add_object(repeated)})}) - if nested - else resources - ) - page[NameObject("/Contents")] = writer._add_object(outer if nested else repeated) - output: Final = io.BytesIO() - writer.write(output) - source: Final = _pdf_source(base64.b64encode(output.getvalue()).decode()) - assert _count_user_content([{"type": "document", "source": source}]) == _count_user_content( - [{"type": "image", "source": source}] - ) - - -def test_pdf_counting_restores_the_callers_parser_limits() -> None: - from pypdf import get_configuration - - original: Final = get_configuration() - source: Final = _pdf_source(_pdf_base64(("Count this",))) - assert _count_user_content([{"type": "document", "source": source}]) > _count_user_content( - [{"type": "image", "source": source}] - ) - assert get_configuration() == original - - -def test_base64_pdf_document_counts_every_page_text_and_rendering(): - """A base64 PDF is read page by page: each page costs its text plus the image Anthropic renders it to. - - Before the fix the whole document was priced as one 85-token image, so a count_tokens call that fell back - to the local counter answered 116 for a 12-page PDF the provider then billed at 35941 input tokens. - """ - prompt: Final = {"type": "text", "text": "Summarize this file."} - first: Final = "Revenue grew eleven percent while churn fell to two percent." - second: Final = "Headcount is flat and the office lease was renewed for three years." - base: Final = _count_user_content([prompt]) - blank_page: Final = ( - _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64(("",)))}]) - base - ) - - assert blank_page > 0 - assert _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64((first, second)))}]) == ( - _count_user_content([prompt, {"type": "text", "text": first}, {"type": "text", "text": second}]) - + 2 * blank_page - ) - one_page: Final = ( - _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64((first,)))}]) - base - ) - twelve_pages: Final = ( - _count_user_content([prompt, {"type": "document", "source": _pdf_source(_pdf_base64((first,) * 12))}]) - base - ) - assert one_page > _count_user_content([prompt, {"type": "text", "text": first}]) - base - assert twelve_pages == 12 * one_page - - -@pytest.mark.parametrize( - "fields", [{}, {"title": "Q3 board packet", "context": "Shared by finance"}], ids=["bare", "described"] -) -def test_pdf_page_rendering_cost_follows_anthropic_image_scaling(fields: dict[str, str]): - """Anthropic rasterizes each PDF page within its image limits (1568 px long edge, 1.15 MP) and bills - width * height / 750 tokens for it: https://platform.claude.com/docs/en/build-with-claude/pdf-support and - https://platform.claude.com/docs/en/build-with-claude/vision, read 2026-10-07, when Bedrock billed about - 1550 tokens per blank Letter page on both Sonnet 4.6 and Opus 4.8. - """ - prompt: Final = {"type": "text", "text": "Summarize this file."} - - def blank_page_cost(width: int, height: int) -> int: - with_page: Final = {"type": "document", "source": _pdf_source(_pdf_base64(("",), width, height)), **fields} - without_page: Final = {"type": "document", "source": _pdf_source(_pdf_base64((), width, height)), **fields} - return _count_user_content([prompt, with_page]) - _count_user_content([prompt, without_page]) - - letter: Final = blank_page_cost(612, 792) - poster: Final = blank_page_cost(2448, 3168) - strip: Final = blank_page_cost(1000, 100) - - assert letter == poster == 1534 - assert strip == 328 - - -def test_pdf_document_without_pypdf_is_priced_like_an_image(monkeypatch: pytest.MonkeyPatch): - prompt: Final = {"type": "text", "text": "Summarize this file."} - source: Final = _pdf_source(_pdf_base64(("Revenue grew eleven percent while churn fell to two percent.",))) - priced_by_page: Final = _count_user_content([prompt, {"type": "document", "source": source}]) - - monkeypatch.setitem(sys.modules, "pypdf", None) - - priced_as_image: Final = _count_user_content([prompt, {"type": "document", "source": source}]) - assert priced_as_image == _count_user_content([prompt, {"type": "image", "source": source}]) - assert priced_as_image < priced_by_page - - def _png_data_url(width: int, height: int) -> str: ihdr = b"\x89PNG\r\n\x1a\n" + (13).to_bytes(4, "big") + b"IHDR" + width.to_bytes(4, "big") + height.to_bytes(4, "big") return "data:image/png;base64," + base64.b64encode(ihdr + b"\x08\x06\x00\x00\x00").decode() diff --git a/uv.lock b/uv.lock index 9cc1507ecf1..fb9d8bbd9ac 100644 --- a/uv.lock +++ b/uv.lock @@ -4825,7 +4825,7 @@ requires-dist = [ { name = "pydantic-settings", specifier = ">=2.14.1,<3.0" }, { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" }, { name = "pynacl", marker = "extra == 'proxy'", specifier = ">=1.6.2,<2.0" }, - { name = "pypdf", marker = "extra == 'proxy-runtime'", specifier = ">=6.19.0,<7.0" }, + { name = "pypdf", marker = "extra == 'proxy-runtime'", specifier = ">=6.16.1,<7.0" }, { name = "pyroscope-io", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.8.16,<1.0" }, { name = "python-dotenv", specifier = ">=1.0.0,<2.0" }, { name = "python-multipart", marker = "extra == 'proxy'", specifier = ">=0.0.27,<1.0" }, From 5e71466355ffe6baa32e7068db1707ba4bded6e5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:52:43 -0700 Subject: [PATCH 35/36] fix(count_tokens): normalize Google system parts and null tools --- litellm/proxy/google_endpoints/endpoints.py | 11 ++++-- .../test_google_gemini_proxy_request.py | 39 ++++++++++++++++++- 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 5fbe780379e..5dbcbcba8e1 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -5,7 +5,7 @@ from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse, ORJSONResponse -from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, field_validator from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types from litellm.proxy._types import * @@ -31,6 +31,11 @@ class _GoogleCountTokensSystemInstruction(BaseModel): model_config = ConfigDict(frozen=True) parts: tuple[_GoogleCountTokensTextPart, ...] + @field_validator("parts", mode="before") + @classmethod + def normalize_parts(cls, parts: object) -> object: + return (parts,) if isinstance(parts, Mapping) else parts + class _GoogleCountTokensFunctionDeclaration(BaseModel): model_config = ConfigDict(frozen=True, populate_by_name=True) @@ -51,7 +56,7 @@ class _GoogleCountTokensRequest(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True) contents: tuple[Mapping[str, JsonValue], ...] = () system_instruction: _GoogleCountTokensSystemInstruction | None = Field(default=None, alias="systemInstruction") - tools: tuple[_GoogleCountTokensTool, ...] = () + tools: tuple[_GoogleCountTokensTool, ...] | None = () router: Final = APIRouter( @@ -111,7 +116,7 @@ def _google_count_tokens_request(model_name: str, data: object) -> TokenCountReq if validated_data.system_instruction is not None else None ) - declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools) + declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools or ()) tools: Final = [ { "name": declaration.name, diff --git a/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py b/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py index 6f8f90efc73..c465b6167e6 100644 --- a/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py +++ b/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py @@ -9,11 +9,12 @@ The request payload is correctly processed and forwarded to the httpx client. import json import unittest.mock -from typing import Optional +from typing import Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from pydantic import JsonValue, ValidationError # Add the parent directory to the system path @@ -493,3 +494,39 @@ async def test_gemini_proxy_config_with_custom_api_base(): if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +@pytest.mark.parametrize("system_field", ("systemInstruction", "system_instruction")) +@pytest.mark.parametrize("parts", ({"text": "Follow these rules"}, [{"text": "Follow these rules"}])) +@pytest.mark.parametrize("tools", (None, [])) +def test_google_count_tokens_accepts_equivalent_system_parts_and_no_tools( + system_field: str, parts: JsonValue, tools: JsonValue +) -> None: + from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request + + contents: Final = [{"role": "user", "parts": [{"text": "Count this request"}]}] + request: Final = _google_count_tokens_request( + "counted", {"contents": contents, system_field: {"parts": parts}, "tools": tools} + ) + + assert request.messages == [{"role": "user", "content": "Count this request"}] + assert request.system == [{"type": "text", "text": "Follow these rules"}] + assert request.tools is None + assert request.contents == contents + + +@pytest.mark.parametrize( + "parts", + ( + {"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}, + [{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}], + ), +) +def test_google_count_tokens_rejects_media_only_system_parts(parts: JsonValue) -> None: + from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request + + with pytest.raises(ValidationError): + _google_count_tokens_request( + "counted", + {"contents": [{"role": "user", "parts": [{"text": "Hello"}]}], "systemInstruction": {"parts": parts}}, + ) From e109f8977ddf3b266d2d6e9376be47374303f479 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:28:56 -0700 Subject: [PATCH 36/36] fix(proxy): restore legacy token count input precedence --- litellm/proxy/proxy_server.py | 12 ------ .../test_vertex_gemini_count_tokens_wire.py | 38 +++++++++++++++- tests/unit/proxy/test_proxy_token_counter.py | 43 ++++--------------- 3 files changed, 45 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 63a920628b6..79b4b8aaa29 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13996,11 +13996,6 @@ _PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( ) -def _drops_params(deployment: Mapping[str, object] | None) -> bool: - litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) - return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) - - def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -14088,13 +14083,6 @@ async def count_request_tokens( model_to_use = _model conversation_empty: Final = not messages and not contents - if prompt and not conversation_empty and not _drops_params(deployment): - raise ProxyException( - message="prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt", - type="invalid_request_error", - param="prompt", - code=400, - ) prompt_only: Final = bool(prompt) and conversation_empty prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES ladder: Final = ( diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 5fe82d37e9e..e54fb1a4a13 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -220,8 +220,8 @@ _CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingPro ) -def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str: - return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key") +def _ai_studio_deployment(scenario: Scenario, api_base: str, **parameters: JsonValue) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key", **parameters) def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: @@ -594,3 +594,37 @@ def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gatew tuple(_json_body(request) for request in requests) == ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9 ) + + +@pytest.mark.parametrize( + ("conversation", "call_endpoint"), + ( + ({"messages": _MESSAGES}, True), + ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]}, True), + ({"messages": _MESSAGES}, False), + ), + ids=("provider-count-messages", "provider-count-contents", "local-count-messages"), +) +@pytest.mark.parametrize("drop_params", (False, True), ids=("no-drop", "drop")) +def test_utils_token_counter_counts_the_conversation_regardless_of_prompt_and_drop_params( + gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool, drop_params: bool +) -> None: + peer: Final = _GeminiPeer() + params: Final = {"call_endpoint": "true" if call_endpoint else "false"} + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=drop_params) + mixed: Final = gateway.request( + "POST", "/utils/token_counter", {"model": model, "prompt": "hi", **conversation}, params=params + ) + mixed_requests: Final = _count_requests(wire) + alone: Final = gateway.request("POST", "/utils/token_counter", {"model": model, **conversation}, params=params) + alone_requests: Final = _count_requests(wire) + assert mixed.status_code == 200, mixed.text + assert alone.status_code == 200, alone.text + mixed_tokens: Final = TypeAdapter(int).validate_python(_JSON_OBJECT.validate_json(mixed.content)["total_tokens"]) + alone_tokens: Final = TypeAdapter(int).validate_python(_JSON_OBJECT.validate_json(alone.content)["total_tokens"]) + assert mixed_tokens == alone_tokens > 0 + assert len(mixed_requests) == len(alone_requests) == int(call_endpoint) + assert [_counted_contents(_json_body(request)) for request in mixed_requests] == [ + _counted_contents(_json_body(request)) for request in alone_requests + ] diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 995089105c4..9703e8ad1dc 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2515,43 +2515,18 @@ _CONVERSATIONS: Final = ( ) -@pytest.mark.asyncio -@pytest.mark.parametrize("conversation", _CONVERSATIONS, ids=["messages", "contents"]) -@pytest.mark.parametrize("call_endpoint", [True, False], ids=["provider", "local"]) -async def test_prompt_next_to_a_conversation_answers_400_without_drop_params( - conversation: dict[str, JsonValue], call_endpoint: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) - client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) - - with ( - patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), - pytest.raises(ProxyException) as rejected, - ): - await token_counter( - request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint - ) - - assert (rejected.value.code, rejected.value.param) == ("400", "prompt") - - @pytest.mark.asyncio @pytest.mark.parametrize( - ("conversation", "call_endpoint", "global_drop_params", "deployment_drop_params"), - [ - (_CONVERSATIONS[0], True, False, True), - (_CONVERSATIONS[1], True, False, True), - (_CONVERSATIONS[0], False, False, True), - (_CONVERSATIONS[0], False, True, False), - ], - ids=[ - "provider_messages", - "provider_contents", - "local_messages", - "local_messages_global_setting", - ], + ("conversation", "call_endpoint"), + ((_CONVERSATIONS[0], True), (_CONVERSATIONS[1], True), (_CONVERSATIONS[0], False)), + ids=("provider_messages", "provider_contents", "local_messages"), ) -async def test_drop_params_counts_a_conversation_without_the_prompt( +@pytest.mark.parametrize( + ("global_drop_params", "deployment_drop_params"), + ((False, False), (False, True), (True, False), (True, True)), + ids=("no_drop", "deployment_drop", "global_drop", "both_drop"), +) +async def test_conversation_takes_precedence_over_prompt_regardless_of_drop_params( conversation: dict[str, JsonValue], call_endpoint: bool, global_drop_params: bool,