mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
fix: count video_url content blocks
This commit is contained in:
parent
1b0ae3af83
commit
b883652907
2 changed files with 78 additions and 1 deletions
|
|
@ -723,6 +723,13 @@ def _count_content_list(
|
|||
num_tokens += _count_image_tokens(
|
||||
image_url, use_default_image_token_count
|
||||
)
|
||||
elif c["type"] == "video_url":
|
||||
video_url = c.get("video_url", "")
|
||||
if isinstance(video_url, dict):
|
||||
video_url = video_url.get("url")
|
||||
if video_url is None:
|
||||
video_url = ""
|
||||
num_tokens += DEFAULT_IMAGE_TOKEN_COUNT + count_function(str(video_url))
|
||||
elif c["type"] in ("tool_use", "tool_result"):
|
||||
num_tokens += _count_anthropic_content(
|
||||
c,
|
||||
|
|
@ -744,7 +751,7 @@ def _count_content_list(
|
|||
)
|
||||
raise ValueError(
|
||||
f"Invalid content item type: {content_type}. "
|
||||
f"Expected str or dict with 'type' field (text, image_url, tool_use, tool_result, thinking)."
|
||||
f"Expected str or dict with 'type' field (text, image_url, video_url, tool_use, tool_result, thinking)."
|
||||
)
|
||||
return num_tokens
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
import litellm
|
||||
from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens
|
||||
from litellm import token_counter as token_counter_old
|
||||
from litellm.constants import DEFAULT_IMAGE_TOKEN_COUNT
|
||||
from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
|
||||
from tests.large_text import text
|
||||
from tests.test_litellm.litellm_core_utils.messages_with_counts import (
|
||||
|
|
@ -908,6 +909,75 @@ def test_token_counter_with_image_url():
|
|||
), f"Expected detail validation error, got: {e}"
|
||||
|
||||
|
||||
def test_token_counter_with_video_url():
|
||||
messages_dict = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Please describe this video."},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {
|
||||
"url": "data:video/mp4;base64,AAAA",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
tokens_dict = token_counter(model="gpt-4o", messages=messages_dict)
|
||||
assert tokens_dict > 0, f"Expected positive token count, got {tokens_dict}"
|
||||
|
||||
messages_str = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": "https://example.com/video.mp4",
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
tokens_str = token_counter(model="gpt-4o", messages=messages_str)
|
||||
assert (
|
||||
tokens_str > 0
|
||||
), f"Expected positive token count for string video_url, got {tokens_str}"
|
||||
assert (
|
||||
tokens_str > DEFAULT_IMAGE_TOKEN_COUNT
|
||||
), f"Expected default video token budget, got {tokens_str}"
|
||||
|
||||
messages_empty_url = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "video_url", "video_url": ""}],
|
||||
}
|
||||
]
|
||||
messages_none_url = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "video_url", "video_url": None}],
|
||||
}
|
||||
]
|
||||
messages_none_nested_url = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "video_url", "video_url": {"url": None}}],
|
||||
}
|
||||
]
|
||||
|
||||
tokens_empty_url = token_counter(model="gpt-4o", messages=messages_empty_url)
|
||||
assert tokens_empty_url > DEFAULT_IMAGE_TOKEN_COUNT
|
||||
assert (
|
||||
token_counter(model="gpt-4o", messages=messages_none_url) == tokens_empty_url
|
||||
)
|
||||
assert (
|
||||
token_counter(model="gpt-4o", messages=messages_none_nested_url)
|
||||
== tokens_empty_url
|
||||
)
|
||||
|
||||
|
||||
def test_token_counter_with_thinking_content():
|
||||
"""
|
||||
Test that _count_content_list() correctly handles Claude's extended thinking content blocks.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue