Merge 8d1489085f into cfc651dcd5

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
@@ -14771,6 +14772,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,
@@ -15339,13 +15358,6 @@ async def count_request_tokens(
     sent_params: Final[Mapping[str, object]] = {**_router_default_litellm_params(llm_router), **request_litellm_params}
     count_deployment: Final = _deployment_with_request_drop_params(deployment, sent_params.get("drop_params"))
     conversation_empty: Final = not messages and not contents
-    if prompt and not conversation_empty and not deployment_drops_params(count_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
     request_messages: Final = _counted_request_messages(messages, messages_format, count_deployment)
     contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents))
@@ -15431,7 +15443,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)

tests/integration/providers/test_vertex_gemini_count_tokens_wire.py

@@ -1043,36 +1043,31 @@ 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"),
 )
-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
+@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()
     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)
+        model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=drop_params)
+        mixed: Final = _count(gateway, {"model": model, "prompt": "hi", **conversation}, call_endpoint)
         mixed_requests: Final = _count_requests(wire)
-        alone: Final = _count(gateway, {"model": model, "messages": _MESSAGES}, call_endpoint)
+        alone: Final = _count(gateway, {"model": model, **conversation}, 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)
+    mixed_tokens: Final = TypeAdapter(int).validate_python(_total_tokens(mixed))
+    alone_tokens: Final = TypeAdapter(int).validate_python(_total_tokens(alone))
+    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
     ]

tests/unit/proxy/test_proxy_token_counter.py

@@ -3424,43 +3424,23 @@ _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),
+        (_CONVERSATIONS[1], False),
+    ),
+    ids=("provider_messages", "provider_contents", "local_messages", "local_contents"),
 )
-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,