diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 76ef3e3b541..450b66745d4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -350,6 +350,7 @@ from litellm.constants import ( WEEKLY_SPEND_REPORT_JOB_ID, ) from litellm.exceptions import RejectedRequestError +from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting @@ -14762,6 +14763,24 @@ def _counted_request_messages( return list(drop_non_string_tool_result_texts(messages, deployment_drops_params(deployment))) +_LOCALLY_COUNTED_MESSAGES: Final = TypeAdapter(list[dict[str, object]]) + + +def _locally_counted_messages( + request_messages: list[dict[str, object]] | None, # mutable-ok: _counted_request_messages returns list messages + contents: Sequence[Mapping[str, object]] | None, + model: str, +) -> list[dict[str, object]] | None: # mutable-ok: token_counter takes list messages + if request_messages or not contents: + return request_messages + translate: Final = cast( # cast-ok: the adapter's untyped **kwargs leave its signature partially unknown + Callable[[str, list[dict[str, object]]], Mapping[str, object]], + GoogleGenAIAdapter().translate_generate_content_to_completion, + ) + translated: Final = translate(model, [dict(content) for content in contents]) + return _LOCALLY_COUNTED_MESSAGES.validate_python(translated.get("messages", [])) + + async def _try_provider_token_count( provider_counter: "BaseTokenCounter", custom_llm_provider: str | None, @@ -15413,7 +15432,19 @@ async def count_request_tokens( 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 request_messages), + [{"role": "user", "content": prompt}] + if prompt_turn + else ( + None + if prompt_only + else _locally_counted_messages( + request_messages, + cast( # cast-ok: TokenCountRequest.contents is typed list[dict], the local counter needs str-keyed parts + Sequence[Mapping[str, object]] | None, contents + ), + model_to_use, + ) + ), ) counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_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 5c08fddd4b8..214bee77895 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -1043,8 +1043,13 @@ def test_utils_token_counter_sends_a_prompt_as_the_user_turn_of_its_messages_for @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"), + ( + ({"messages": _MESSAGES}, True), + ({"contents": _CONTENTS}, True), + ({"messages": _MESSAGES}, False), + ({"contents": _CONTENTS}, False), + ), + ids=("provider-count-messages", "provider-count-contents", "local-count-messages", "local-count-contents"), ) @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( diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 436694c3849..be7e50adf62 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -3422,8 +3422,13 @@ _CONVERSATIONS: Final = ( @pytest.mark.asyncio @pytest.mark.parametrize( ("conversation", "call_endpoint"), - ((_CONVERSATIONS[0], True), (_CONVERSATIONS[1], True), (_CONVERSATIONS[0], False)), - ids=("provider_messages", "provider_contents", "local_messages"), + ( + (_CONVERSATIONS[0], True), + (_CONVERSATIONS[1], True), + (_CONVERSATIONS[0], False), + (_CONVERSATIONS[1], False), + ), + ids=("provider_messages", "provider_contents", "local_messages", "local_contents"), ) @pytest.mark.parametrize( ("global_drop_params", "deployment_drop_params"),