mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(proxy): restore legacy token count input precedence
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
21ab2cb87b
commit
2367019447
3 changed files with 19 additions and 61 deletions
|
|
@ -15328,13 +15328,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))
|
||||
|
|
|
|||
|
|
@ -1046,33 +1046,23 @@ def test_utils_token_counter_sends_a_prompt_as_the_user_turn_of_its_messages_for
|
|||
(({"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
|
||||
@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
|
||||
]
|
||||
|
|
|
|||
|
|
@ -3419,43 +3419,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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue