fix(proxy): count non-text contents parts in the local token_counter fallback

contents carrying only media or function-call parts produced no chat
messages, so the local fallback raised ValueError -> 500; serialize
each part to its JSON frame with base64 blobs elided
This commit is contained in:
shrey kharbanda 2026-09-24 04:51:02 +00:00
parent bb90256ba1
commit 3e1ae630ff
2 changed files with 74 additions and 14 deletions

View file

@ -13459,24 +13459,57 @@ def _system_message(system: object) -> ChatCompletionSystemMessage | None:
return message
def _elide_binary_blobs(part: object) -> object:
"""Replace base64 blobs (inlineData.data, file data) so serialized
parts stay a sane size for the local tokenizer."""
if isinstance(part, Mapping):
return {
key: ("<binary>" if key == "data" and isinstance(value, str) else _elide_binary_blobs(value))
for key, value in part.items()
}
if isinstance(part, list):
return [_elide_binary_blobs(item) for item in part]
return part
def _part_to_text(part: object) -> str:
if isinstance(part, Mapping) and isinstance(part.get("text"), str):
return part["text"]
return json.dumps(_elide_binary_blobs(part), default=str)
def _content_parts(content: Mapping[str, object]) -> tuple[object, ...]:
parts: Final = content.get("parts")
if isinstance(parts, list):
return tuple(parts)
return (content,)
def _contents_as_messages(contents: object) -> tuple[Mapping[str, object], ...] | None:
"""Approximate gemini contents as chat messages for the local fallback
tokenizer; only text parts are countable locally."""
if not isinstance(contents, list):
tokenizer. Text parts count as text; other parts count as their JSON
frame with base64 blobs elided."""
if contents is None:
return None
messages: Final = tuple(
if isinstance(contents, list):
messages: Final = tuple(
{ # mutable-ok: transient chat-shaped message for the local tokenizer
"role": "assistant" if content.get("role") == "model" else "user",
"content": "\n".join(_part_to_text(part) for part in _content_parts(content)),
}
for content in contents
if isinstance(content, Mapping)
)
counted: Final = tuple(message for message in messages if message["content"])
if counted:
return counted
fallback: Final[tuple[Mapping[str, object], ...]] = (
{ # mutable-ok: transient chat-shaped message for the local tokenizer
"role": "assistant" if content.get("role") == "model" else "user",
"content": "\n".join(
part["text"]
for part in content.get("parts", ())
if isinstance(part, Mapping) and isinstance(part.get("text"), str)
),
}
for content in contents
if isinstance(content, Mapping)
"role": "user",
"content": json.dumps(_elide_binary_blobs(contents), default=str),
},
)
return tuple(message for message in messages if message["content"]) or None
return fallback
@router.post(

View file

@ -423,5 +423,32 @@ def test_token_counter_contents_only_request_counts_text_parts(client, auth_as,
assert response.status_code == 200, response.text
assert response.json()["total_tokens"] == litellm.token_counter(
model="claude-fable-5",
messages=[{"role": "user", "content": "hello world"}, {"role": "assistant", "content": "hi there"}],
messages=[
{"role": "user", "content": 'hello world\n{"inline_data": {"data": "<binary>"}}'},
{"role": "assistant", "content": "hi there"},
{"role": "user", "content": '{"inline_data": {"data": "<binary>"}}'},
],
)
def test_token_counter_media_only_contents_falls_back_instead_of_500(client, auth_as, monkeypatch):
"""Regression: contents carrying only media or function-call parts (no text)
still produce a local count, because their JSON frames serialize to text."""
monkeypatch.setattr(proxy_server, "llm_router", None)
monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False)
contents = [
{"role": "user", "parts": [{"inline_data": {"mime_type": "image/png", "data": "aGVsbG8="}}]},
{"role": "model", "parts": [{"function_call": {"name": "get_weather", "args": {"city": "sf"}}}]},
]
with auth_as():
response = client.post("/utils/token_counter", json={"model": "claude-fable-5", "contents": contents})
assert response.status_code == 200, response.text
assert response.json()["total_tokens"] == litellm.token_counter(
model="claude-fable-5",
messages=[
{"role": "user", "content": '{"inline_data": {"mime_type": "image/png", "data": "<binary>"}}'},
{"role": "assistant", "content": '{"function_call": {"name": "get_weather", "args": {"city": "sf"}}}'},
],
)