mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
fix(proxy)!: count prompt-only token_counter requests with the provider when call_endpoint is set (#46839)
* fix(proxy): count prompt-only token_counter requests with the provider when call_endpoint is set * fix(proxy): keep the local count for a prompt-only token_counter request when the provider count fails * fix(proxy): log a prompt-only token count fallback as a one-line warning * fix(proxy): send a prompt-only token_counter request to the provider only on Gemini and Vertex AI * fix(proxy): count a token_counter prompt sent next to an empty messages or contents list * fix(proxy): count a prompt next to messages on the local token count instead of answering 500 * fix(proxy): answer 400 for a token count prompt next to a conversation unless drop_params is set With drop_params the conversation is counted without the prompt. On the local count, a prompt sent with system or tools is counted as a user message with them, like the messages form. * test(proxy): count tools with a prompt in the token counter fallback test * test(proxy): cover prompt counting on /utils/token_counter in the Gemini count wire tests * test(proxy): name the token counter fallback test by what it checks
This commit is contained in:
Mateo Wang
GitHub
parent
642a2ad934
commit
2dc6aafa6b
4 changed files with 421 additions and 13 deletions
|
|
@ -14488,6 +14488,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
|
||||
|
|
@ -14554,8 +14564,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
|
||||
|
|
@ -14566,13 +14576,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,
|
||||
|
|
@ -14605,8 +14625,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)
|
||||
|
|
@ -14616,7 +14638,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,
|
||||
|
|
|
|||
|
|
@ -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 == ()
|
||||
|
|
|
|||
|
|
@ -390,8 +390,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"
|
||||
|
|
@ -412,4 +411,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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue