fix(proxy): only attach tools to the count_tokens fallback when counting messages

This commit is contained in:
mateo-berri 2026-08-28 06:10:30 -07:00
parent 70ba0bb973
commit 83ab87091b
2 changed files with 31 additions and 1 deletions

View file

@ -12244,7 +12244,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages)
)
counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats
list[ChatCompletionToolParam] | None, tools
list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None
)
total_tokens: Final = await asyncify(litellm.token_counter)(
model=model_to_use,

View file

@ -229,3 +229,33 @@ def test_token_counter_fallback_counts_tools_system_and_anthropic_blocks(client,
tools=tools,
)
assert full > bare
def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, monkeypatch):
"""
Regression: a raw-text ``prompt`` request that also carries ``tools`` (no ``messages``) must
still count. ``litellm.token_counter`` rejects tools on the text path, so the fallback route
only attaches tools when it is counting messages; otherwise this 500'd instead of returning
the plain text count.
"""
monkeypatch.setattr(proxy_server, "llm_router", None)
monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False)
prompt = "count the tokens in this sentence please"
tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Look up the current weather for a city",
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
},
}
]
with auth_as():
response = client.post(
"/utils/token_counter", json={"model": "claude-fable-5", "prompt": prompt, "tools": tools}
)
assert response.status_code == 200, response.text
assert response.json()["total_tokens"] == litellm.token_counter(model="claude-fable-5", text=prompt)