fix: handle image_url=None in token_counter to prevent ValueError crash

This commit is contained in:
Dantuluri Surya Narayana Raju 2026-05-17 17:27:17 +05:30
parent cf9b5e4fa7
commit 7400c1b1fa
2 changed files with 36 additions and 1 deletions

View file

@ -585,7 +585,8 @@ def _count_image_tokens(
Count tokens for an image_url content block.
Args:
image_url: The image URL data - can be a string URL or dict with 'url' and 'detail'
image_url: The image URL data - can be a string URL or dict with 'url' and 'detail'.
None is treated as an unknown image and returns the default token count.
use_default_image_token_count: Whether to use default image token counts
Returns:
@ -594,6 +595,8 @@ def _count_image_tokens(
Raises:
ValueError: If image_url is invalid type or detail value is invalid
"""
if image_url is None:
return DEFAULT_IMAGE_TOKEN_COUNT
if isinstance(image_url, dict):
detail = image_url.get("detail", "auto")
if detail not in ["low", "high", "auto"]:

View file

@ -908,6 +908,38 @@ def test_token_counter_with_image_url():
), f"Expected detail validation error, got: {e}"
def test_token_counter_image_url_none():
"""
Regression test: token_counter must not raise ValueError when a content
block has {"type": "image_url", "image_url": None}.
Previously _count_image_tokens() reached its final else-branch and raised:
ValueError: Invalid image_url type: NoneType. Expected str or dict with 'url' field.
After the fix, None is treated as an unknown image and the default token
count (DEFAULT_IMAGE_TOKEN_COUNT = 250) is returned so callers are not
disrupted by a None value in an otherwise valid message list.
"""
from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "hello"},
{"type": "image_url", "image_url": None},
],
}
]
tokens = token_counter(model="gpt-4-vision-preview", messages=messages)
assert tokens > 0, f"Expected positive token count, got {tokens}"
assert tokens >= DEFAULT_IMAGE_TOKEN_COUNT, (
f"Expected at least DEFAULT_IMAGE_TOKEN_COUNT={DEFAULT_IMAGE_TOKEN_COUNT} "
f"tokens due to the None image_url fallback, got {tokens}"
)
def test_token_counter_with_thinking_content():
"""
Test that _count_content_list() correctly handles Claude's extended thinking content blocks.