mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
fix(token_counter): count video/audio blocks as 0 instead of raising
`_count_content_list` raised `Invalid content item type` for `video_url` and `input_audio` blocks, so any request carrying video or audio failed in the counter. Their payload is opaque here, so they contribute 0 and the enumeration in the catch-all error now lists them. The catch-all moves into a helper so the new case does not raise the complexity of `_count_content_list`. On the ollama_chat route the counter runs after generation, so this turned a response Ollama had already produced into a 500. Fixes #28071
This commit is contained in:
parent
68c0a972a7
commit
eb32c8cccd
3 changed files with 136 additions and 7 deletions
|
|
@ -897,6 +897,26 @@ def _count_anthropic_content(
|
|||
return tokens
|
||||
|
||||
|
||||
def _count_media_block(block: object) -> int:
|
||||
"""
|
||||
Count a block type the other branches do not handle.
|
||||
|
||||
Video and audio blocks carry an opaque data URI whose cost the provider decides
|
||||
(frame sampling rate, audio chunking), so they contribute 0 rather than raise - on
|
||||
the ollama route the counter runs after generation, where raising discards a
|
||||
response the model already produced. Any other type is unknown and raises.
|
||||
"""
|
||||
if isinstance(block, dict) and block.get("type") in ("video_url", "input_audio"):
|
||||
return 0
|
||||
content_type = block.get("type", type(block).__name__) if isinstance(block, dict) else type(block).__name__
|
||||
raise ValueError(
|
||||
f"Invalid content item type: {content_type}. "
|
||||
f"Expected str or dict with 'type' field "
|
||||
f"(text, image_url, image, document, file, video_url, input_audio, tool_use, tool_result, "
|
||||
f"thinking, redacted_thinking, tool_reference)."
|
||||
)
|
||||
|
||||
|
||||
def _count_content_list(
|
||||
count_function: TokenCounterFunction,
|
||||
content_list: str
|
||||
|
|
@ -966,13 +986,7 @@ def _count_content_list(
|
|||
if tool_name:
|
||||
num_tokens += count_function(tool_name)
|
||||
else:
|
||||
content_type = c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__
|
||||
raise ValueError(
|
||||
f"Invalid content item type: {content_type}. "
|
||||
f"Expected str or dict with 'type' field "
|
||||
f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, "
|
||||
f"tool_reference)."
|
||||
)
|
||||
num_tokens += _count_media_block(c)
|
||||
return num_tokens
|
||||
except Exception as e:
|
||||
if default_token_count is not None:
|
||||
|
|
|
|||
|
|
@ -1830,3 +1830,51 @@ def test_messages_reach_token_count_honours_disable_token_counter(monkeypatch: p
|
|||
messages = _threshold_test_messages(turns=3)
|
||||
assert messages_reach_token_count(model="gpt-4o", messages=messages, threshold=0) is True
|
||||
assert messages_reach_token_count(model="gpt-4o", messages=messages, threshold=1) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"content_block",
|
||||
[
|
||||
{"type": "video_url", "video_url": {"url": "data:video/mp4;base64,AAAA"}},
|
||||
{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "mp3"}},
|
||||
],
|
||||
ids=["video_url", "input_audio"],
|
||||
)
|
||||
def test_token_counter_with_non_text_modality_blocks(content_block: dict):
|
||||
"""
|
||||
Video and audio content blocks must not raise.
|
||||
|
||||
Raising `Invalid content item type: video_url` breaks two paths: on the proxy
|
||||
it nulls response_cost and drops the SpendLogs row, and on the ollama route the
|
||||
counter runs after generation, so an already-produced 200 response becomes a 500.
|
||||
|
||||
Their payload is opaque here, so they contribute 0 tokens and the count matches
|
||||
the same message without the block.
|
||||
"""
|
||||
prompt = {"type": "text", "text": "Describe the attachment."}
|
||||
messages = [{"role": "user", "content": [prompt, content_block]}]
|
||||
text_only = [{"role": "user", "content": [prompt]}]
|
||||
|
||||
tokens = token_counter_new(model="gpt-4o", messages=messages)
|
||||
assert tokens > 0, f"Expected positive token count, got {tokens}"
|
||||
assert tokens == token_counter_new(model="gpt-4o", messages=text_only)
|
||||
|
||||
|
||||
def test_count_content_list_error_message_lists_modality_types():
|
||||
"""
|
||||
The catch-all error enumerates the handled block types, and the video and audio
|
||||
types must appear there.
|
||||
"""
|
||||
from litellm.litellm_core_utils.token_counter import _count_content_list
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid content item type: totally_unknown_block") as exc_info:
|
||||
_count_content_list(
|
||||
count_function=len,
|
||||
content_list=[{"type": "totally_unknown_block"}],
|
||||
use_default_image_token_count=False,
|
||||
default_token_count=None,
|
||||
)
|
||||
|
||||
message = str(exc_info.value)
|
||||
for content_type in ("video_url", "input_audio"):
|
||||
assert content_type in message
|
||||
|
|
|
|||
|
|
@ -989,3 +989,70 @@ class TestOllamaStreamingUsage:
|
|||
)
|
||||
|
||||
assert result.usage is None
|
||||
|
||||
|
||||
class TestOllamaChatVideoBlockUsage:
|
||||
"""Usage counting for requests that carry a non-text content block."""
|
||||
|
||||
@staticmethod
|
||||
def _transform(
|
||||
config: OllamaChatConfig,
|
||||
ollama_response: dict,
|
||||
messages: list[AllMessageValues],
|
||||
) -> ModelResponse:
|
||||
import json
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.types.utils import Choices, Message, ModelResponse
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = ollama_response
|
||||
mock_response.text = json.dumps(ollama_response)
|
||||
|
||||
model_response = ModelResponse()
|
||||
model_response.choices = [Choices(message=Message(content=""), index=0)]
|
||||
|
||||
return config.transform_response(
|
||||
model="qwen3:14b",
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
messages=messages,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
api_key=None,
|
||||
json_mode=False,
|
||||
)
|
||||
|
||||
def test_missing_counts_with_a_video_block_do_not_raise(self):
|
||||
"""
|
||||
End-to-end of the two fixes on this route: Ollama omits the counts, so the
|
||||
estimator runs over a message that carries a `video_url` block. It must
|
||||
produce usage instead of raising, because the response has already been
|
||||
generated by the time the counter runs.
|
||||
"""
|
||||
ollama_response = {
|
||||
"model": "qwen3:14b",
|
||||
"created_at": "2025-01-11T00:00:00.000000Z",
|
||||
"message": {"role": "assistant", "content": "A short clip."},
|
||||
"done": True,
|
||||
}
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this."},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": "data:video/mp4;base64,AAAA"},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
result = self._transform(OllamaChatConfig(), ollama_response, messages)
|
||||
|
||||
assert result.usage.prompt_tokens > 0
|
||||
assert result.usage.completion_tokens > 0
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue