scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-07 07:58:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 150

Merge 41a857c34e into 7e37c20748

Browse source
This commit is contained in:
John Karam 2026-10-30 21:21:58 +00:00 • committed by GitHub
commit f192f681e9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 343 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

50
litellm/litellm_core_utils/token_counter.py
View file

@ -47,11 +47,14 @@ from litellm.constants import (
TOKEN_COUNTER_MAX_EXACT_CHARS,
)
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.core_helpers import as_str_mapping, has_attributes
from litellm.litellm_core_utils.prompt_templates.chat_template_tools import (
ChatTemplateToolSection,
chat_template_tool_section,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_custom_tool_format_to_responses_shape,
)
from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding
if TYPE_CHECKING:
@ -80,7 +83,7 @@ from litellm.types.llms.openai import (
ChatCompletionToolParam,
OpenAIMessageContentListBlock,
)
from litellm.types.utils import Message, SelectTokenizerResponse
from litellm.types.utils import Message, SelectTokenizerResponse, is_custom_tool_call_dict
class UnsupportedContentTypeError(ValueError):
@ -890,10 +893,11 @@ def _count_function_call_tokens(
raise ValueError(f"Unsupported type {type(value)} for key tool_calls in message {message}")
total = 0
for tool_call in value:
if "function" not in tool_call:
raise ValueError(f"Unsupported tool call {tool_call} must contain a function key")
function_arguments = tool_call["function"].get("arguments", "")
total += count_function(str(function_arguments))
if (fields := _tool_call_fields(tool_call)) is None or (
"function" not in fields and "custom" not in fields
):
raise ValueError(f"Unsupported tool call {tool_call} must contain a function or custom key")
total += count_function(_tool_call_text(fields))
return total
if key == "function_call":
if not isinstance(value, Mapping):
@ -902,6 +906,24 @@ def _count_function_call_tokens(
raise ValueError(f"Unexpected key {key!r}; expected 'tool_calls' or 'function_call'")
class _ToolCallFields(Protocol):
def get(self, key: str, default: object = None, /) -> object: ...
def __contains__(self, key: object, /) -> bool: ...
def _tool_call_fields(value: object) -> _ToolCallFields | None:
return value if has_attributes(value, _ToolCallFields, "get", "__contains__") else None
def _tool_call_text(tool_call: _ToolCallFields) -> str:
is_custom: Final = is_custom_tool_call_dict({"type": tool_call.get("type"), "custom": tool_call.get("custom")})
payload: Final = _tool_call_fields(tool_call.get("custom" if is_custom else "function"))
if payload is None:
return ""
return str(payload.get("input" if is_custom else "arguments", ""))
def _count_messages(
params: _MessageCountParams,
messages: list[AllMessageValues],
@ -1801,6 +1823,10 @@ def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mappi
if isinstance(function, Mapping):
yield function
return
custom: Final = as_str_mapping(tool.get("custom"))
if "input_schema" not in tool and (custom is not None or tool.get("type") == "custom"):
yield _custom_tool_definition(custom if custom is not None else tool)
return
declarations: Final = tool.get("function_declarations") or tool.get("functionDeclarations")
if isinstance(declarations, list):
for declaration in declarations:
@ -1825,6 +1851,18 @@ def _function_definitions_for_tool(tool: Mapping[str, object]) -> Iterable[Mappi
}
def _custom_tool_definition(custom: Mapping[str, object]) -> Mapping[str, object]:
format_obj: Final = as_str_mapping(custom.get("format"))
grammar: Final = (
convert_custom_tool_format_to_responses_shape(format_obj).get("definition") if format_obj is not None else None
)
description_parts: Final = tuple(str(part) for part in (custom.get("description"), grammar) if part)
return {
"name": custom.get("name"),
**({"description": "\n".join(description_parts)} if description_parts else {}),
}
def _format_single_function_definition(function: Mapping[str, object]) -> tuple[str, ...]:
function_name: Final = function.get("name")
if not function_name:

79
litellm/llms/openai/responses/count_tokens/transformation.py
View file

@ -5,14 +5,17 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp
"""
import copy
from collections.abc import Mapping, Sequence
from collections.abc import Iterable, Mapping, Sequence
from typing import Final, Literal, Protocol, TypeAlias
from openai.types.responses.input_token_count_params import InputTokenCountParams
from pydantic import JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items, require_attributes
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_custom_tool_format_to_responses_shape,
)
from litellm.litellm_core_utils.tool_fields import DropPolicy
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.anthropic.messages.convertible_blocks import malformed_tool_result_text_error
@ -23,7 +26,7 @@ from litellm.llms.anthropic.messages.responses_adapters.transformation import (
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema
from litellm.types.llms.openai import ResponsesAPIRequestParams, replayable_responses_input
from litellm.types.utils import all_litellm_params
from litellm.types.utils import all_litellm_params, is_custom_tool_call_dict
class ResponsesInputTextPart(TypedDict):
@ -82,6 +85,47 @@ def _function_fields(value: object) -> _FunctionFields:
return require_attributes(value, _FunctionFields, "get", "__contains__", "__getitem__")
def _is_custom_tool_call(tool_call: _FieldReader) -> bool:
return is_custom_tool_call_dict({"type": tool_call.get("type"), "custom": tool_call.get("custom")})
def _assistant_tool_calls(messages: Sequence[Mapping[str, object]]) -> Iterable[_FieldReader]:
for msg in messages:
tool_calls = msg.get("tool_calls")
if msg.get("role") == "assistant" and tool_calls:
yield from (_field_reader(tool_call) for tool_call in iter_items(tool_calls))
def _responses_tool_call_item(tool_call: _FieldReader) -> Mapping[str, object]:
call_id: Final = tool_call.get("id", "")
if _is_custom_tool_call(tool_call):
custom: Final = _field_reader(tool_call.get("custom") or {})
return {
"type": "custom_tool_call",
"call_id": call_id,
"name": custom.get("name", ""),
"input": custom.get("input", ""),
}
function: Final = _field_reader(tool_call.get("function") or {})
return {
"type": "function_call",
"call_id": call_id,
"name": function.get("name", ""),
"arguments": function.get("arguments", ""),
}
def _responses_custom_tool(custom: Mapping[str, object]) -> Mapping[str, object]:
description: Final = custom.get("description")
format_obj: Final = as_str_mapping(custom.get("format"))
return {
"type": "custom",
"name": custom.get("name", ""),
**({"description": description} if description is not None else {}),
**({"format": convert_custom_tool_format_to_responses_shape(format_obj)} if format_obj is not None else {}),
}
_ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema)
@ -249,6 +293,8 @@ class OpenAICountTokensConfig:
"strict": bool(tool.get("strict")),
}
)
elif tool.get("type") == "custom" and (custom := as_str_mapping(tool.get("custom"))) is not None:
transformed.append(dict(_responses_custom_tool(custom)))
else:
# Pass through non-function tools (e.g., web_search, file_search)
transformed.append(tool)
@ -267,6 +313,9 @@ class OpenAICountTokensConfig:
"""
input_items: Final[list[dict[str, object]]] = []
instructions_parts: Final[list[str]] = []
custom_tool_call_ids: Final = frozenset(
tool_call.get("id") for tool_call in _assistant_tool_calls(messages) if _is_custom_tool_call(tool_call)
)
for msg in messages:
role = msg.get("role", "")
@ -297,26 +346,20 @@ class OpenAICountTokensConfig:
if content:
input_items.append({"role": "assistant", "content": content})
if tool_calls:
for tc in iter_items(tool_calls):
tool_call = _field_reader(tc)
function = tool_call.get("function", {})
call_id = tool_call.get("id", "")
function_fields = _field_reader(function)
input_items.append(
{
"type": "function_call",
"call_id": call_id,
"name": function_fields.get("name", ""),
"arguments": function_fields.get("arguments", ""),
}
)
input_items.extend(
dict(_responses_tool_call_item(_field_reader(tool_call)))
for tool_call in iter_items(tool_calls)
)
elif not content:
input_items.append({"role": "assistant", "content": content})
elif role == "tool":
tool_call_id = msg.get("tool_call_id", "")
input_items.append(
{
"type": "function_call_output",
"call_id": msg.get("tool_call_id", ""),
"type": "custom_tool_call_output"
if tool_call_id in custom_tool_call_ids
else "function_call_output",
"call_id": tool_call_id,
"output": content if isinstance(content, str) else str(content),
}
)

149
tests/unit/litellm_core_utils/test_token_counter.py
View file

@ -50,7 +50,13 @@ from litellm.litellm_core_utils.token_counter import (
from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from litellm.types.utils import SelectTokenizerResponse
from litellm.types.utils import (
ChatCompletionCustomToolCallPayload,
ChatCompletionMessageCustomToolCall,
ChatCompletionMessageToolCall,
Function,
SelectTokenizerResponse,
)
from tests.large_text import text
from tests.unit.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
@ -385,6 +391,146 @@ def test_token_counter_legacy_function_call_counts_arguments():
assert fc_tokens > 500, f"4000-char arguments payload must contribute real tokens, got {fc_tokens}"
def _tool_call_turn(tool_call: object) -> list[dict[str, object]]:
return [
{"role": "user", "content": "run the code"},
{"role": "assistant", "content": None, "tool_calls": [tool_call]},
{"role": "tool", "tool_call_id": "call_1", "content": "done"},
]
_CUSTOM_CODE = "print('hello world from a custom tool')\n" * 50
@pytest.mark.parametrize(
"custom_call",
(
{"id": "call_1", "type": "custom", "custom": {"name": "code_exec", "input": _CUSTOM_CODE}},
{"id": "call_1", "type": "custom", "function": None, "custom": {"name": "code_exec", "input": _CUSTOM_CODE}},
{"id": "call_1", "custom": {"name": "code_exec", "input": _CUSTOM_CODE}},
),
ids=("typed", "sdk-shaped-null-function", "typeless"),
)
def test_token_counter_custom_tool_call_counts_input_like_function_arguments(custom_call: dict[str, object]):
custom_tokens = token_counter(model="gpt-4o", messages=_tool_call_turn(custom_call))
function_tokens = token_counter(
model="gpt-4o",
messages=_tool_call_turn(
{"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": _CUSTOM_CODE}}
),
)
assert custom_tokens == function_tokens, (
f"custom tool call input must count like function arguments, got {custom_tokens} vs {function_tokens}"
)
assert custom_tokens > 300, f"custom tool call input must contribute real tokens, got {custom_tokens}"
@pytest.mark.parametrize(
("sdk_call", "dict_call"),
(
(
ChatCompletionMessageToolCall(
id="call_1", type="function", function=Function(name="code_exec", arguments=_CUSTOM_CODE)
),
{"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": _CUSTOM_CODE}},
),
(
ChatCompletionMessageCustomToolCall(
id="call_1", custom=ChatCompletionCustomToolCallPayload(name="code_exec", input=_CUSTOM_CODE)
),
{"id": "call_1", "type": "custom", "custom": {"name": "code_exec", "input": _CUSTOM_CODE}},
),
(
{"id": "call_1", "type": "function", "function": Function(name="code_exec", arguments=_CUSTOM_CODE)},
{"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": _CUSTOM_CODE}},
),
),
ids=("function", "custom", "dict-holding-function-object"),
)
def test_token_counter_counts_reused_sdk_tool_call_objects_like_their_dicts(
sdk_call: object, dict_call: dict[str, object]
):
sdk_tokens = token_counter(model="gpt-4o", messages=_tool_call_turn(sdk_call))
assert sdk_tokens == token_counter(model="gpt-4o", messages=_tool_call_turn(dict_call))
assert sdk_tokens > 300, f"a reused SDK tool call must contribute its payload tokens, got {sdk_tokens}"
@pytest.mark.parametrize(
("null_call", "empty_call"),
(
(
{"id": "call_1", "type": "custom", "custom": None},
{"id": "call_1", "type": "custom", "custom": {"name": "code_exec", "input": ""}},
),
(
{"id": "call_1", "type": "function", "function": None},
{"id": "call_1", "type": "function", "function": {"name": "code_exec", "arguments": ""}},
),
),
ids=("custom", "function"),
)
def test_token_counter_null_tool_call_payload_counts_as_empty_call(
null_call: dict[str, object], empty_call: dict[str, object]
):
assert token_counter(model="gpt-4o", messages=_tool_call_turn(null_call)) == token_counter(
model="gpt-4o", messages=_tool_call_turn(empty_call)
)
def test_token_counter_rejects_tool_call_without_function_or_custom_payload():
with pytest.raises(ValueError, match="must contain a function or custom key"):
token_counter(model="gpt-4o", messages=_tool_call_turn({"id": "call_1", "type": "function"}))
def test_token_counter_counts_custom_tool_definition_description_and_grammar():
messages = [{"role": "user", "content": "run the code"}]
grammar = "start: " + " | ".join(f'"command_{index}"' for index in range(100))
def custom_tool(**fields: object) -> dict[str, object]:
return {"type": "custom", "custom": {"name": "code_exec", **fields}}
no_tools = token_counter(model="gpt-4o", messages=messages)
bare = token_counter(model="gpt-4o", messages=messages, tools=[custom_tool()])
described = token_counter(
model="gpt-4o",
messages=messages,
tools=[
custom_tool(
description="Runs one shell command",
format={"type": "grammar", "grammar": {"syntax": "lark", "definition": grammar}},
)
],
)
assert bare > no_tools, f"a custom tool's name must count, got {bare} vs {no_tools} without tools"
assert described - bare > 300, f"description and grammar must count, got {described} vs {bare} bare"
def test_token_counter_counts_responses_shaped_custom_tool_like_the_chat_shape():
messages = [{"role": "user", "content": "run the code"}]
grammar = "start: " + " | ".join(f'"command_{index}"' for index in range(100))
chat_tool = {
"type": "custom",
"custom": {
"name": "code_exec",
"description": "Runs one shell command",
"format": {"type": "grammar", "grammar": {"syntax": "lark", "definition": grammar}},
},
}
responses_tool = {
"type": "custom",
"name": "code_exec",
"description": "Runs one shell command",
"format": {"type": "grammar", "syntax": "lark", "definition": grammar},
}
assert token_counter(model="gpt-4o", messages=messages, tools=[responses_tool]) == token_counter(
model="gpt-4o", messages=messages, tools=[chat_tool]
)
@pytest.mark.parametrize(
"message_count_pair",
MESSAGES_TEXT,
@ -744,6 +890,7 @@ def test_token_counter_counts_anthropic_tool_choices_like_openai_choices(
{**_anthropic_tool(_WEATHER_FUNCTION), "strict": True},
{**_anthropic_tool(_WEATHER_FUNCTION), "defer_loading": True, "input_examples": [{"location": "Paris"}]},
{**_anthropic_tool(_WEATHER_FUNCTION), "custom": {"eager_input_streaming": True}},
{**_anthropic_tool(_WEATHER_FUNCTION), "type": "custom", "custom": {"eager_input_streaming": True}},
_anthropic_tool(_WEB_SEARCH_FUNCTION),
],
)

91
tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py
View file

@ -110,8 +110,24 @@ _SCHEMA: Final[dict[str, object]] = {"type": "object", "properties": {}} # muta
"strict": False,
},
),
(
{
"type": "custom",
"name": "lookup",
"description": "Find a record",
"input_schema": _SCHEMA,
"custom": {"eager_input_streaming": True},
},
{
"type": "function",
"name": "lookup",
"description": "Find a record",
"parameters": _SCHEMA,
"strict": False,
},
),
),
ids=("anthropic", "anthropic-strict", "anthropic-custom-type"),
ids=("anthropic", "anthropic-strict", "anthropic-custom-type", "anthropic-custom-options"),
)
def test_transform_anthropic_tool_shapes(tool: dict[str, object], expected: dict[str, object]) -> None:
result: Final = OpenAICountTokensConfig().transform_request_to_count_tokens(
@ -227,6 +243,79 @@ def test_messages_to_responses_input_reads_a_tool_call_that_forwards_its_fields(
}
def test_messages_to_responses_input_maps_custom_tool_calls_and_their_outputs():
messages: Final = [
{"role": "user", "content": "Run it"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{"id": "call_custom", "type": "custom", "custom": {"name": "code_exec", "input": "print(2+3)"}},
{"id": "call_fn", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}},
],
},
{"role": "tool", "tool_call_id": "call_custom", "content": "5"},
{"role": "tool", "tool_call_id": "call_fn", "content": "sunny"},
]
input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages)
assert input_items[1:] == [
{"type": "custom_tool_call", "call_id": "call_custom", "name": "code_exec", "input": "print(2+3)"},
{"type": "function_call", "call_id": "call_fn", "name": "get_weather", "arguments": "{}"},
{"type": "custom_tool_call_output", "call_id": "call_custom", "output": "5"},
{"type": "function_call_output", "call_id": "call_fn", "output": "sunny"},
]
@pytest.mark.parametrize(
("tool_call", "expected"),
(
(
{"id": "call_1", "type": "custom", "custom": None},
{"type": "custom_tool_call", "call_id": "call_1", "name": "", "input": ""},
),
(
{"id": "call_1", "type": "function", "function": None},
{"type": "function_call", "call_id": "call_1", "name": "", "arguments": ""},
),
),
ids=("custom", "function"),
)
def test_messages_to_responses_input_maps_null_tool_call_payloads_to_empty_calls(
tool_call: dict[str, object], expected: dict[str, object]
) -> None:
input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(
[{"role": "user", "content": "Run it"}, {"role": "assistant", "content": None, "tool_calls": [tool_call]}]
)
assert input_items[1:] == [expected]
def test_transform_custom_tool_uses_the_flat_responses_shape():
tool: Final = {
"type": "custom",
"custom": {
"name": "code_exec",
"description": "Run python",
"format": {"type": "grammar", "grammar": {"syntax": "lark", "definition": "start: /.+/"}},
},
}
result: Final = OpenAICountTokensConfig().transform_request_to_count_tokens(
model="gpt-4o", input="Run it", tools=[tool]
)
assert result["tools"] == [
{
"type": "custom",
"name": "code_exec",
"description": "Run python",
"format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"},
}
]
def test_messages_to_responses_input_preserves_images():
"""An image block must survive the round trip, or OpenAI counts only the text.