mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(count_tokens): preserve image inputs when counting Responses API tokens
The chat-to-Responses reverse transform kept only text blocks, so an image input was dropped before the count went to OpenAI. A 256x256 image request counted 13 tokens instead of 268.
This commit is contained in:
parent
6b7159323b
commit
73ab647b1c
3 changed files with 169 additions and 10 deletions
|
|
@ -4,7 +4,69 @@ OpenAI Responses API token counting transformation logic.
|
|||
This module handles the transformation of requests to OpenAI's /v1/responses/input_tokens endpoint.
|
||||
"""
|
||||
|
||||
from typing import Any, Final
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
|
||||
class ResponsesInputTextPart(TypedDict):
|
||||
type: ReadOnly[Literal["input_text"]]
|
||||
text: ReadOnly[str]
|
||||
|
||||
|
||||
class ResponsesInputImagePart(TypedDict):
|
||||
type: ReadOnly[Literal["input_image"]]
|
||||
image_url: ReadOnly[str]
|
||||
detail: ReadOnly[str]
|
||||
|
||||
|
||||
ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart
|
||||
|
||||
|
||||
def _chat_image_block_to_responses_part(image_url: object) -> ResponsesInputImagePart | None:
|
||||
url: Final = image_url.get("url") if isinstance(image_url, Mapping) else image_url
|
||||
if not isinstance(url, str) or not url:
|
||||
return None
|
||||
detail: Final = image_url.get("detail") if isinstance(image_url, Mapping) else None
|
||||
part: Final[ResponsesInputImagePart] = {
|
||||
"type": "input_image",
|
||||
"image_url": url,
|
||||
"detail": detail if isinstance(detail, str) and detail else "auto",
|
||||
}
|
||||
return part
|
||||
|
||||
|
||||
def _chat_block_to_responses_part(block: object) -> ResponsesInputPart | None:
|
||||
if isinstance(block, str):
|
||||
bare: Final[ResponsesInputTextPart] = {"type": "input_text", "text": block}
|
||||
return bare
|
||||
if not isinstance(block, Mapping):
|
||||
return None
|
||||
match block.get("type"):
|
||||
case "text":
|
||||
text_value: Final = block.get("text")
|
||||
text: Final[ResponsesInputTextPart] = {
|
||||
"type": "input_text",
|
||||
"text": text_value if isinstance(text_value, str) else "",
|
||||
}
|
||||
return text
|
||||
case "image_url":
|
||||
return _chat_image_block_to_responses_part(block.get("image_url"))
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def chat_content_blocks_to_responses_content(
|
||||
content: Sequence[object],
|
||||
) -> str | tuple[ResponsesInputPart, ...]:
|
||||
"""Text-only content collapses to a joined string, so text-only counts stay unchanged."""
|
||||
parts: Final = tuple(
|
||||
part for part in (_chat_block_to_responses_part(block) for block in content) if part is not None
|
||||
)
|
||||
if any(part["type"] != "input_text" for part in parts):
|
||||
return parts
|
||||
return "\n".join(part["text"] for part in parts if part["type"] == "input_text")
|
||||
|
||||
|
||||
class OpenAICountTokensConfig:
|
||||
|
|
@ -120,14 +182,7 @@ class OpenAICountTokensConfig:
|
|||
instructions_parts.append("\n".join(text_parts))
|
||||
elif role == "user":
|
||||
if isinstance(content, list):
|
||||
# Extract text from content blocks for Responses API
|
||||
text_parts = []
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get("type") == "text":
|
||||
text_parts.append(block.get("text", ""))
|
||||
elif isinstance(block, str):
|
||||
text_parts.append(block)
|
||||
content = "\n".join(text_parts)
|
||||
content = chat_content_blocks_to_responses_content(content)
|
||||
input_items.append({"role": "user", "content": content})
|
||||
elif role == "assistant":
|
||||
# Map tool_calls to Responses API function_call items
|
||||
|
|
|
|||
|
|
@ -163,6 +163,103 @@ def test_messages_to_responses_input_with_tool():
|
|||
}
|
||||
|
||||
|
||||
def test_messages_to_responses_input_preserves_images():
|
||||
"""An image block must survive the round trip, or OpenAI counts only the text.
|
||||
|
||||
A 256x256 image is worth 255 tokens to OpenAI's counting API; dropping it
|
||||
turned a 268-token request into a 13-token one.
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is in this image?"},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "data:image/png;base64,iVBORw0KGgo=", "detail": "high"},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
|
||||
|
||||
assert instructions is None
|
||||
assert input_items == [
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
{"type": "input_text", "text": "What is in this image?"},
|
||||
{
|
||||
"type": "input_image",
|
||||
"image_url": "data:image/png;base64,iVBORw0KGgo=",
|
||||
"detail": "high",
|
||||
},
|
||||
),
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_messages_to_responses_input_image_without_detail_defaults_to_auto():
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}],
|
||||
}
|
||||
]
|
||||
|
||||
input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages)
|
||||
|
||||
assert input_items[0]["content"] == (
|
||||
{"type": "input_image", "image_url": "https://example.com/cat.png", "detail": "auto"},
|
||||
)
|
||||
|
||||
|
||||
def test_messages_to_responses_input_bare_string_image_url_is_preserved():
|
||||
messages = [{"role": "user", "content": [{"type": "image_url", "image_url": "https://example.com/cat.png"}]}]
|
||||
|
||||
input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages)
|
||||
|
||||
assert input_items[0]["content"] == (
|
||||
{"type": "input_image", "image_url": "https://example.com/cat.png", "detail": "auto"},
|
||||
)
|
||||
|
||||
|
||||
def test_messages_to_responses_input_text_only_blocks_stay_a_joined_string():
|
||||
"""Text-only content must keep collapsing to a string so existing counts do not shift."""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}],
|
||||
}
|
||||
]
|
||||
|
||||
input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages)
|
||||
|
||||
assert input_items == [{"role": "user", "content": "first\nsecond"}]
|
||||
|
||||
|
||||
def test_messages_to_responses_input_drops_unmappable_blocks():
|
||||
"""A block with no Responses API equivalent is skipped, never forwarded verbatim."""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "hi"},
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}},
|
||||
{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages)
|
||||
|
||||
assert input_items[0]["content"] == (
|
||||
{"type": "input_text", "text": "hi"},
|
||||
{"type": "input_image", "image_url": "https://example.com/cat.png", "detail": "auto"},
|
||||
)
|
||||
|
||||
|
||||
def test_validate_request_valid():
|
||||
"""Test that valid requests pass validation."""
|
||||
config = OpenAICountTokensConfig()
|
||||
|
|
|
|||
|
|
@ -3,10 +3,12 @@ Test for response_api_endpoints/endpoints.py
|
|||
"""
|
||||
|
||||
import unittest
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from httpx import Response
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
|
@ -1816,7 +1818,12 @@ class TestResponsesInputTokens:
|
|||
never registered, so the POST fell through to the GET/DELETE-only
|
||||
/v1/responses/{response_id} route and returned 405."""
|
||||
|
||||
def _post_input_tokens(self, body, path="/v1/responses/input_tokens", counter=None):
|
||||
def _post_input_tokens(
|
||||
self,
|
||||
body: dict[str, Any],
|
||||
path: str = "/v1/responses/input_tokens",
|
||||
counter: AsyncMock | None = None,
|
||||
) -> tuple[Response, AsyncMock]:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.response_api_endpoints.endpoints import _proxy_token_counter
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue