From 35f52bf3e83d9b8d7c4d58fd784ee1347b732620 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:33:50 -0700 Subject: [PATCH 01/13] fix(responses): fold bridged Responses output into one chat choice (#46099) * fix(responses): fold bridged Responses output into one chat choice A Chat Completions request bridged to /responses got one choice per Responses message item, so a reply shaped [message, message, function_call] put the tool call on choices[1] and clients reading only choices[0] never saw it. Fold every message text, annotation, reasoning item and tool call into choices[0] Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): drop mutable annotations and the extra cast in the choice fold Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): remove a cast-ok marker that suppressed nothing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): join folded message texts with a newline Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * revert(responses): join folded message texts with no separator so stream and non-stream match Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit 4b7b8efd4cf91c302ecb3196bc7ac3a4288ebd51) --- .../transformation.py | 110 +++++--- .../test_responses_bridge_incomplete.py | 256 +++++++++++++++++- ...responses_transformation_transformation.py | 60 +++- 3 files changed, 384 insertions(+), 42 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 41f6a211036..5ad9f854313 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -74,7 +74,7 @@ if TYPE_CHECKING: ChatCompletionToolReferenceObject, OpenAIMessageContentListBlock, ) - from litellm.types.utils import Choices + from litellm.types.utils import Choices, Message _CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage")) @@ -308,6 +308,67 @@ def _as_chat_reasoning_items( return cast(list[ChatCompletionReasoningItem], list(reasoning_items)) +def _shift_annotation(annotation: ChatCompletionAnnotation, offset: int) -> ChatCompletionAnnotation: + shifted: Final = { + key: value + offset if key in ("start_index", "end_index") and isinstance(value, int) else value + for key, value in annotation.items() + } + return cast(ChatCompletionAnnotation, shifted) + + +def _message_annotations(message: "Message") -> Sequence[ChatCompletionAnnotation]: + annotations: Sequence[ChatCompletionAnnotation] | None = getattr(message, "annotations", None) + return annotations or () + + +def _message_reasoning_content(message: "Message") -> str | None: + reasoning_content: str | None = getattr(message, "reasoning_content", None) + return reasoning_content + + +def _message_reasoning_items(message: "Message") -> Sequence[ChatCompletionReasoningItem]: + reasoning_items: Sequence[ChatCompletionReasoningItem] | None = getattr(message, "reasoning_items", None) + return reasoning_items or () + + +def _shifted_annotations(messages: Sequence["Message"], offsets: Sequence[int]) -> Iterator[ChatCompletionAnnotation]: + for message, offset in zip(messages, offsets): + for annotation in _message_annotations(message): + yield _shift_annotation(annotation, offset) + + +def _fold_choices_into_one( + choices: Sequence["Choices"], + tool_calls: Sequence[Mapping[str, object]], + trailing_reasoning_content: str | None, + trailing_reasoning_item: _BuiltReasoningItem | None, +) -> "Choices": + """Clients that only read choices[0] must still see every text and tool call.""" + from litellm.types.utils import Choices, Message + + messages: Final = tuple(choice.message for choice in choices) + texts: Final = tuple(str(message.content or "") for message in messages) + offsets: Final = (0, *accumulate(len(text) for text in texts)) + annotations: Final = tuple(_shifted_annotations(messages, offsets)) + reasoning_contents: Final = (*(_message_reasoning_content(m) for m in messages), trailing_reasoning_content) + reasoning_items: Final = ( + *chain.from_iterable(_message_reasoning_items(m) for m in messages), + *(() if trailing_reasoning_item is None else (trailing_reasoning_item,)), + ) + return Choices( + message=Message( + role=messages[0].role if messages else "assistant", + content="".join(texts) if messages else None, + annotations=list(annotations) or None, + tool_calls=list(tool_calls) or None, + reasoning_content=" ".join(value for value in reasoning_contents if value) or None, + reasoning_items=_as_chat_reasoning_items(reasoning_items), + ), + finish_reason="tool_calls" if tool_calls else "stop", + index=0, + ) + + _ToolChoiceT = TypeVar("_ToolChoiceT") @@ -951,47 +1012,16 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): else: pass # don't fail request if item in list is not supported - if accumulated_tool_calls and choices: - last_choice: Final = choices[-1] - last_reasoning_content: Final = getattr(last_choice.message, "reasoning_content", None) - last_reasoning_items: Final = getattr(last_choice.message, "reasoning_items", None) - merged_reasoning_content: Final = ( - " ".join(value for value in (last_reasoning_content, reasoning_content) if value) or None - ) - merged_reasoning_items: Final = _as_chat_reasoning_items( - ( - *(last_reasoning_items or ()), - *(() if pending_reasoning_item is None else (pending_reasoning_item,)), - ) - ) - merged_message: Final = Message( - role=last_choice.message.role, - content=last_choice.message.content, - annotations=getattr(last_choice.message, "annotations", None), + if not choices and not accumulated_tool_calls: + return choices + return [ + _fold_choices_into_one( + choices=choices, tool_calls=accumulated_tool_calls, - reasoning_content=merged_reasoning_content, - reasoning_items=merged_reasoning_items, + trailing_reasoning_content=reasoning_content, + trailing_reasoning_item=pending_reasoning_item, ) - return [ - *choices[:-1], - Choices(message=merged_message, finish_reason="tool_calls", index=last_choice.index), - ] - - if accumulated_tool_calls: - msg = Message( - content=None, - tool_calls=accumulated_tool_calls, - reasoning_content=reasoning_content, - reasoning_items=cast( - list[ChatCompletionReasoningItem] | None, - ([pending_reasoning_item] if pending_reasoning_item is not None else None), - ), - ) - choices.append(Choices(message=msg, finish_reason="tool_calls", index=index)) - reasoning_content = None - pending_reasoning_item = None - - return choices + ] @staticmethod def _build_empty_incomplete_choice( diff --git a/tests/integration/providers/test_responses_bridge_incomplete.py b/tests/integration/providers/test_responses_bridge_incomplete.py index 5d3877c954e..6e005f928ef 100644 --- a/tests/integration/providers/test_responses_bridge_incomplete.py +++ b/tests/integration/providers/test_responses_bridge_incomplete.py @@ -3,8 +3,9 @@ import uuid from typing import Final import pytest -from integration._support.client import Gateway +from integration._support.client import Gateway, list_value, object_value from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue @pytest.mark.covers("other.provider_wire.responses_bridge.max_output_tokens_incomplete_maps_to_length") @@ -572,3 +573,256 @@ def test_chat_over_responses_deployment_merges_function_call_followed_by_message } ], response.text assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +_WEATHER_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} + + +def _message_item(item_id: str, text: str) -> dict[str, JsonValue]: + return { + "type": "message", + "id": item_id, + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + + +_WEATHER_CALL: Final[dict[str, JsonValue]] = { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + "status": "completed", +} + +_TWO_MESSAGES_THEN_CALL: Final = ( + _message_item("msg_first", "First message."), + _message_item("msg_second", "Second message."), + _WEATHER_CALL, +) + + +def test_chat_over_responses_deployment_folds_two_messages_and_function_call_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-two-messages-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + body=json.dumps( + { + "id": "resp_two_messages", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-6-sol", + "output": list(_TWO_MESSAGES_THEN_CALL), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [_WEATHER_TOOL], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["choices"] == [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": "First message.Second message.", + "tool_calls": [ + { + "id": "fc_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + "index": 0, + } + ], + }, + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + + +def _responses_frame(event: str, payload: dict[str, JsonValue]) -> bytes: + return f"event: {event}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n".encode() + + +def _message_item_frames(output_index: int, item: dict[str, JsonValue]) -> tuple[bytes, ...]: + item_id: Final = item["id"] + text: Final = object_value(list_value(item["content"])[0])["text"] + return ( + _responses_frame( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": {**item, "status": "in_progress", "content": []}, + }, + ), + _responses_frame( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": item_id, + "output_index": output_index, + "content_index": 0, + "delta": text, + }, + ), + _responses_frame( + "response.output_item.done", + {"type": "response.output_item.done", "output_index": output_index, "item": item}, + ), + ) + + +def _function_call_frames(output_index: int, item: dict[str, JsonValue]) -> tuple[bytes, ...]: + return ( + _responses_frame( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": {**item, "status": "in_progress", "arguments": ""}, + }, + ), + _responses_frame( + "response.function_call_arguments.delta", + { + "type": "response.function_call_arguments.delta", + "item_id": item["id"], + "output_index": output_index, + "delta": item["arguments"], + }, + ), + _responses_frame( + "response.function_call_arguments.done", + { + "type": "response.function_call_arguments.done", + "item_id": item["id"], + "output_index": output_index, + "arguments": item["arguments"], + }, + ), + _responses_frame( + "response.output_item.done", + {"type": "response.output_item.done", "output_index": output_index, "item": item}, + ), + ) + + +def _sse_payloads(lines: list[str]) -> list[dict[str, JsonValue]]: + return [json.loads(line[6:]) for line in lines if line.startswith("data: ") and line != "data: [DONE]"] + + +def test_chat_stream_over_responses_deployment_folds_two_messages_and_function_call_into_one_choice( + gateway: Gateway, +) -> None: + identity: Final = "responses-bridge-two-messages-stream-" + uuid.uuid4().hex + response_payload: Final[dict[str, JsonValue]] = { + "id": "resp_two_messages_stream", + "object": "response", + "created_at": 1789788253, + "model": "gpt-6-sol", + "status": "in_progress", + "output": [], + } + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') + assert request.method == "POST" and request.target == "/responses", request.target + assert json.loads(request.body)["stream"] is True, request.body + return Reply( + content_type="text/event-stream", + chunks=( + _responses_frame("response.created", {"type": "response.created", "response": response_payload}), + *_message_item_frames(0, _TWO_MESSAGES_THEN_CALL[0]), + *_message_item_frames(1, _TWO_MESSAGES_THEN_CALL[1]), + *_function_call_frames(2, _TWO_MESSAGES_THEN_CALL[2]), + _responses_frame( + "response.completed", + { + "type": "response.completed", + "response": { + **response_payload, + "status": "completed", + "output": list(_TWO_MESSAGES_THEN_CALL), + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + }, + ), + ), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}], + "tools": [_WEATHER_TOOL], + "stream": True, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + chunks: Final = _sse_payloads(list(response.iter_lines())) + choices: Final = [object_value(choice) for chunk in chunks for choice in list_value(chunk["choices"])] + assert [choice["index"] for choice in choices] == [0] * len(choices), chunks + deltas: Final = [object_value(choice["delta"]) for choice in choices] + assert ( + "".join(str(delta["content"]) for delta in deltas if delta.get("content")) + == "First message.Second message." + ) + tool_calls: Final = [ + object_value(call) + for delta in deltas + if delta.get("tool_calls") + for call in list_value(delta["tool_calls"]) + ] + assert [ + object_value(call["function"])["name"] for call in tool_calls if object_value(call["function"]).get("name") + ] == ["get_weather"], chunks + assert ( + "".join(str(object_value(call["function"]).get("arguments") or "") for call in tool_calls) + == '{"city":"Paris"}' + ), chunks + assert [choice["finish_reason"] for choice in choices if choice.get("finish_reason")] == ["tool_calls"], chunks + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 8c6c4059a22..863504dcd29 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -5019,4 +5019,62 @@ def test_transform_request_drop_params_in_litellm_params_gates_the_prompt_cache_ litellm_logging_obj=Mock(), ) - assert "prompt_cache_breakpoint" not in result["input"][0]["content"][0] \ No newline at end of file + assert "prompt_cache_breakpoint" not in result["input"][0]["content"][0] + + +def test_convert_response_output_folds_two_messages_and_function_call_into_one_choice() -> None: + first: Final = ResponseOutputMessage.model_validate( + { + "id": "msg_first", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "First message.", "annotations": []}], + } + ) + second_annotation: Final = { + "type": "url_citation", + "start_index": 0, + "end_index": 6, + "title": "Second", + "url": "https://example.com/second", + } + second: Final = ResponseOutputMessage.model_validate( + { + "id": "msg_second", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Second message.", "annotations": [second_annotation]}], + } + ) + function_call: Final = ResponseFunctionToolCall( + id="fc_1", + type="function_call", + status="completed", + arguments='{"city":"Paris"}', + call_id="call_1", + name="get_weather", + ) + + choices: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( # pyright: ignore[reportPrivateUsage] # exercises the converter directly + (first, second, function_call) + ) + + assert len(choices) == 1 + assert choices[0].index == 0 + assert choices[0].finish_reason == "tool_calls" + assert choices[0].message.content == "First message.Second message." + assert choices[0].message.annotations == [ + {**second_annotation, "start_index": len("First message."), "end_index": len("First message.") + 6} + ] + tool_calls: Final = choices[0].message.tool_calls + assert tool_calls is not None + assert [call.model_dump(exclude_none=True) for call in tool_calls] == [ + { + "id": "fc_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + "index": 0, + } + ] From 1c6e82a60e949edb34199d502e7f7528683a22cb Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 13 Oct 2026 14:39:52 -0700 Subject: [PATCH 02/13] fix(mcp): resolve a bare tool name on a single-server scope against that server (#46225) A JSON-RPC tools/call on /{server_name}/mcp (or /mcp with x-mcp-servers naming one server) that uses the upstream's bare tool name answered "Tool not found" on every proxy process that had not yet listed that server, because the per-process tool-name map stays empty for a server whose upstream only lists tools with the caller's own credentials, and mcp SDK 2 no longer lists before an uncached call. call_mcp_tool now hands the single scoped server to execute_mcp_tool as requested_server_id, so the existing REST contract lists the server with the caller's headers and routes the bare name to it, and a bare name shared by two servers goes to the one the path names Co-authored-by: aayushbaluni <73417844+aayushbaluni@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit ec28acf322cf48f32a5c0c26524c961be351a117) --- .../_experimental/mcp_server/operations.py | 16 ++++ .../mcp_server/test_mcp_server.py | 41 ++++------ .../mcp_server/test_operations.py | 82 +++++++++++++++++++ 3 files changed, 112 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 2f0b8e64c9e..af236843a1f 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -2450,6 +2450,20 @@ async def fire_mcp_tool_call_failure_logging( ) +def _single_scoped_server( + mcp_servers: Sequence[str] | None, allowed_mcp_servers: Sequence[MCPServer] +) -> MCPServer | None: + """The one server a ``/{server_name}/mcp`` path or ``x-mcp-servers`` header scoped this call to. + + A bare tool name on such a request resolves against that server, the way a REST ``server_id`` + does, instead of against the per-process tool-name map, which stays empty for a server whose + upstream only lists tools with the caller's own credentials until this process has listed it. + """ + if mcp_servers is None or len(mcp_servers) != 1 or len(allowed_mcp_servers) != 1: + return None + return allowed_mcp_servers[0] + + @client @catalog_operation(lambda: global_mcp_server_manager) async def call_mcp_tool( @@ -2511,6 +2525,7 @@ async def call_mcp_tool( detail="User not allowed to call this tool.", ) + scoped_server: Final = _single_scoped_server(mcp_servers, allowed_mcp_servers) # Delegate to execute_mcp_tool for execution response = await execute_mcp_tool( name=name, @@ -2527,6 +2542,7 @@ async def call_mcp_tool( input_responses=input_responses, request_state=request_state, incoming_bearer_token=incoming_bearer_token, + requested_server_id=scoped_server.server_id if scoped_server is not None else None, **kwargs, ) except Exception as e: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py index 449e437865b..f8ffcd6c31e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py @@ -2994,11 +2994,7 @@ async def test_call_mcp_tool_uses_manager_permission_lookup(): "get_mcp_server_by_id", return_value=mock_server, ), - patch.object( - global_mcp_server_manager, - "get_mcp_server_from_tool_name", - return_value=mock_server, - ) as mock_get_server, + patch.object(global_mcp_server_manager, "server_exposes_tool", return_value=True), patch( "litellm.proxy._experimental.mcp_server.operations.global_mcp_tool_registry" ) as mock_tool_registry, @@ -3023,24 +3019,18 @@ async def test_call_mcp_tool_uses_manager_permission_lookup(): assert result == expected_response mock_get_allowed.assert_awaited_once() - # We call `_get_mcp_server_from_tool_name` multiple times: - # - for logging/metadata - # - for resolving the server during dispatch - # - and inside `call_tool` for guardrails/hooks - # The exact count isn't important, only that it is used. - assert mock_get_server.call_count >= 2 - # First call should use the prefixed tool name - assert ( - mock_get_server.call_args_list[0][0][0] - == f"{mock_server.name}/gmail_send_email" - ) + # The request names exactly one server, so that server is the dispatch target + # and the tool name reaches it untouched (the "/" is not a LiteLLM prefix). + mock_handle_managed.assert_awaited_once() + assert mock_handle_managed.await_args.kwargs["server_name"] == mock_server.name + assert mock_handle_managed.await_args.kwargs["name"] == f"{mock_server.name}/gmail_send_email" @pytest.mark.asyncio async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permissions(): """ - Ensure `call_mcp_tool` correctly resolves the MCP server for an unprefixed tool - name and enforces server-level permissions using that resolved server. + Ensure `call_mcp_tool` resolves an unprefixed tool name on a request scoped to + one server against that server and enforces server-level permissions with it. """ from litellm.proxy._experimental.mcp_server.server import ( call_mcp_tool, @@ -3070,11 +3060,7 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission "get_mcp_server_by_id", return_value=mock_server, ), - patch.object( - global_mcp_server_manager, - "get_mcp_server_from_tool_name", - return_value=mock_server, - ) as mock_get_server, + patch.object(global_mcp_server_manager, "server_exposes_tool", return_value=True), patch( "litellm.proxy._experimental.mcp_server.operations.global_mcp_tool_registry" ) as mock_tool_registry, @@ -3091,7 +3077,7 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission mock_tool_registry.get_tool.return_value = None mock_handle_managed.return_value = expected_response - # Call with UNPREFIXED tool name; server should be resolved via mapping + # Call with UNPREFIXED tool name; the single server the request names is the target result = await call_mcp_tool( name="gmail_send_email", arguments={"body": "hello"}, @@ -3100,8 +3086,9 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission assert result == expected_response mock_get_allowed.assert_awaited_once() - # We should resolve the server at least once using the unprefixed name - assert mock_get_server.call_count >= 1 - assert mock_get_server.call_args_list[0][0][0] == "gmail_send_email" + mock_handle_managed.assert_awaited_once() + assert mock_handle_managed.await_args.kwargs["server_name"] == mock_server.name + assert mock_handle_managed.await_args.kwargs["name"] == "gmail_send_email" # Permissions check should be invoked with the resolved server name mock_is_allowed.assert_called_once() + assert mock_is_allowed.call_args.kwargs["server_name"] == mock_server.name diff --git a/tests/unit/proxy/_experimental/mcp_server/test_operations.py b/tests/unit/proxy/_experimental/mcp_server/test_operations.py index dcc4eb1a570..af547ebf6c7 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_operations.py @@ -565,6 +565,88 @@ async def test_tools_call_warmup_does_not_consume_mcp_server_rpm() -> None: assert [tool.name for tool in listing.tools] == ["echo"] + +def _scoped_rig( + monkeypatch: pytest.MonkeyPatch, +) -> tuple[MCPServer, MCPServer, AsyncMock, AsyncMock]: + """Two config servers that both expose a bare ``list_rules`` on a manager that has listed neither.""" + from mcp.types import CallToolResult, TextContent + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + manager: Final = MCPServerManager() + memory: Final = MCPServer( + server_id="id-memory", name="memory_server", server_name="memory_server", transport=MCPTransport.http + ) + other: Final = MCPServer( + server_id="id-other", name="other_server", server_name="other_server", transport=MCPTransport.http + ) + manager.config_mcp_servers = {server.server_id: server for server in (memory, other)} + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + monkeypatch.setattr( + manager, "get_allowed_mcp_servers", AsyncMock(return_value=[memory.server_id, other.server_id]) + ) + monkeypatch.setattr(global_mcp_tool_registry, "published_tools", {}) + listing: Final = AsyncMock() + monkeypatch.setattr(operations, "_get_tools_from_mcp_servers", listing) + handled: Final = AsyncMock( + return_value=CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + ) + monkeypatch.setattr(operations, "_handle_managed_mcp_tool", handled) + return memory, other, listing, handled + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", ["memory_server", "other_server"]) +async def test_bare_tool_name_on_a_single_server_scope_routes_to_that_server_cold( + monkeypatch: pytest.MonkeyPatch, scope: str +) -> None: + """``/{server_name}/mcp`` (or ``x-mcp-servers`` naming one server) with the upstream's bare tool name + must list and call the scoped server on a process whose tool-name map has never seen it, and when two + servers share the bare name the scope decides (BerriAI/litellm#44831).""" + memory, other, listing, handled = _scoped_rig(monkeypatch) + expected: Final = memory if scope == "memory_server" else other + + async with operations.global_mcp_server_manager.catalog.operation(): + result: Final = await operations.call_mcp_tool( + name="list_rules", + arguments={}, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-scoped"), + mcp_servers=[scope], + ) + + assert result.is_error is False + assert listing.await_args.kwargs["mcp_servers"] == [expected.server_id] + assert handled.await_args.kwargs["server_name"] == expected.name + assert handled.await_args.kwargs["name"] == "list_rules" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mcp_servers", [None, ["memory_server", "other_server"]]) +async def test_bare_tool_name_without_a_single_server_scope_is_still_unresolved_cold( + monkeypatch: pytest.MonkeyPatch, mcp_servers: list[str] | None +) -> None: + """The scope shortcut only applies when the request names exactly one server; an ambiguous bare name + keeps today's answer rather than being guessed onto the first allowed server.""" + from fastapi import HTTPException + + _memory, _other, listing, handled = _scoped_rig(monkeypatch) + + async with operations.global_mcp_server_manager.catalog.operation(): + with pytest.raises(HTTPException) as unresolved: + await operations.call_mcp_tool( + name="list_rules", + arguments={}, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-scoped"), + mcp_servers=mcp_servers, + ) + + assert unresolved.value.status_code == 404 + listing.assert_not_awaited() + handled.assert_not_awaited() + @pytest.mark.asyncio async def test_tools_call_pre_call_check_enforces_mcp_server_rpm() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager From cbe5c288b8b203de6a5ffa56335cd7ef0d5677af Mon Sep 17 00:00:00 2001 From: Chloe Lu <3273594+clonylu@users.noreply.github.com> Date: Thu, 15 Oct 2026 14:30:03 +0800 Subject: [PATCH 03/13] fix(anthropic): keep signed empty thinking blocks on /v1/messages requests (#41362) The request-side sanitizer on /v1/messages dropped every thinking block whose text is empty, even with a valid signature. Claude Opus 4.8 / Fable 5 / Fable 5.1 return exactly that shape by default (thinking.display omitted): the signature is the preserved reasoning. Replaying the assistant turn as returned therefore lost the model's reasoning on every multi-turn request, silently (the provider answers 200 without it), and left preserved-thinking enforcement (thinking.block_binding) nothing to bind. Anthropic accepts an empty thinking block that carries a signature and binds it to its prefix; only an empty block with no signature (the LIT-6357 case, a turn a non-Anthropic model served through the bridge) is rejected with "each thinking block must contain thinking". Use the signature-aware predicate is_empty_unsigned_thinking_block, which #38809 already applied on the emit side, for the request side too. The unsigned case is still dropped. Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit 897bfea38d22666c67be8e8e523ae316f4e42265) --- litellm/llms/anthropic/common_utils.py | 42 +++++++++++-------- .../anthropic/test_anthropic_common_utils.py | 40 +++++++++++++++--- 2 files changed, 59 insertions(+), 23 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 16de0e119c2..d037f76ed03 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1678,7 +1678,7 @@ def strip_empty_content_blocks_from_anthropic_messages( ) -> list[Any]: """ Return a new message list with empty or whitespace-only ``{"type": "text"}`` - and ``{"type": "thinking"}`` content blocks removed. + content blocks and empty *unsigned* ``{"type": "thinking"}`` blocks removed. Anthropic's API rejects requests containing such blocks with ``"messages: text content blocks must be non-empty"`` and @@ -1694,8 +1694,13 @@ def strip_empty_content_blocks_from_anthropic_messages( on the unified ``/v1/messages`` path. ``/v1/chat/completions`` already handles this in ``anthropic_messages_pt``; this helper provides the equivalent guarantee for the native Anthropic Messages path. - ``redacted_thinking`` blocks are never touched: they carry opaque - ``data`` instead of thinking text. + + A thinking block with empty text but a non-empty ``signature`` is kept: + it is the default response shape of Claude Opus 4.8 / Fable 5 / Fable 5.1 + (``thinking.display`` omitted), the signature is the model's preserved + reasoning, and Anthropic accepts it on input. ``redacted_thinking`` + blocks are never touched: they carry opaque ``data`` instead of thinking + text. Messages whose content is a list and becomes empty after stripping are omitted, matching :func:`strip_thinking_blocks_from_anthropic_messages`. @@ -1708,7 +1713,7 @@ def strip_empty_content_blocks_from_anthropic_messages( out.append(m) continue content = m["content"] - filtered = [b for b in content if not _is_empty_text_block(b) and not is_empty_thinking_block(b)] + filtered = [b for b in content if not _is_empty_text_block(b) and not is_empty_unsigned_thinking_block(b)] if len(filtered) == len(content): out.append(m) elif filtered: @@ -1726,11 +1731,11 @@ def _is_empty_text_block(block: object) -> bool: def is_empty_thinking_block(block: object) -> bool: """ True for a ``{"type": "thinking"}`` content block whose thinking text is - missing, not a string, or empty/whitespace-only after ``.strip()``. - Anthropic rejects such blocks with ``"each thinking block must contain - thinking"`` (whitespace-only included, verified live), regardless of any - signature they carry. ``redacted_thinking`` blocks are a different type - and always return False. + missing, not a string, or empty/whitespace-only after ``.strip()``, + whatever its signature. This is a shape check only; whether such a block + may be sent to Anthropic depends on the signature, see + :func:`is_empty_unsigned_thinking_block`. ``redacted_thinking`` blocks are + a different type and always return False. """ if not isinstance(block, dict) or block.get("type") != "thinking": return False @@ -1742,15 +1747,16 @@ def is_empty_unsigned_thinking_block(block: object) -> bool: """ True for an empty ``{"type": "thinking"}`` block carrying no signature. - The emit-side predicate: response paths drop a thinking block only when it - holds nothing the client could need. A signature-only block is a real - provider response (Bedrock Converse under adaptive thinking emits a - reasoning block with empty text and only a signature) and the client needs - the signature to replay reasoning across tool-use turns, so it must be - emitted. Request paths keep using :func:`is_empty_thinking_block`: - Anthropic rejects empty thinking blocks in request history regardless of - signature, and the inbound strip self-heals a replayed signature-only - block. + Used on both the emit and the request side: a thinking block is dropped + only when it holds nothing the provider or the client could use. A + signature-only block is a real provider response (Claude Opus 4.8 / + Fable 5 / Fable 5.1 return it by default, and Bedrock Converse under + adaptive thinking emits a reasoning block with empty text and only a + signature); the client needs the signature to replay reasoning across + tool-use turns, and Anthropic accepts the replayed block and binds it to + its prefix. Only an empty block with no signature (a turn a + non-Anthropic reasoning model served through the bridge) is rejected + with ``"each thinking block must contain thinking"`` and must go. """ if not isinstance(block, dict) or not is_empty_thinking_block(block): return False diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index 17840a829c1..ce6f9754c74 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -1422,11 +1422,11 @@ class TestAnthropicThinkingSignatureSelfHeal: assert is_empty_unsigned_thinking_block({"type": "redacted_thinking", "data": "opaque"}) is False assert is_empty_unsigned_thinking_block("not a dict") is False - def test_strip_empty_content_blocks_drops_empty_thinking_blocks(self): + def test_strip_empty_content_blocks_drops_empty_unsigned_thinking_blocks(self): """LIT-6357 ingestion half: an assistant tool-loop turn carrying an - empty (even signed) thinking block keeps its tool_use blocks and loses - the poison; whitespace-only counts as empty; a non-empty thinking block - and redacted_thinking are untouched.""" + empty, unsigned thinking block keeps its tool_use blocks and loses the + poison; whitespace-only counts as empty; a non-empty thinking block and + redacted_thinking are untouched.""" from litellm.llms.anthropic.common_utils import ( strip_empty_content_blocks_from_anthropic_messages, ) @@ -1436,7 +1436,7 @@ class TestAnthropicThinkingSignatureSelfHeal: {"role": "user", "content": "weather?"}, { "role": "assistant", - "content": [{"type": "thinking", "thinking": "", "signature": "sig_abc"}, tu], + "content": [{"type": "thinking", "thinking": "", "signature": ""}, tu], }, { "role": "assistant", @@ -1455,6 +1455,36 @@ class TestAnthropicThinkingSignatureSelfHeal: assert out[2]["content"][0]["thinking"] == "real plan" assert len(msgs[1]["content"]) == 2 + def test_strip_empty_content_blocks_keeps_signed_empty_thinking_blocks(self): + """A thinking block with empty text and a signature is the default + response shape of Claude Opus 4.8 / Fable 5 / Fable 5.1 (thinking.display + omitted): the signature IS the preserved reasoning, Anthropic accepts it + on input and binds it to its prefix. Dropping it silently discards the + model's reasoning on every replayed turn and leaves preserved-thinking + enforcement (thinking.block_binding) nothing to check. The message and + block order must survive untouched.""" + from litellm.llms.anthropic.common_utils import ( + strip_empty_content_blocks_from_anthropic_messages, + ) + + tu = {"type": "tool_use", "id": "toolu_01A", "name": "get_weather", "input": {"city": "Paris"}} + signed = {"type": "thinking", "thinking": "", "signature": "EqQBCkYIBxgCKkD"} + msgs = [ + {"role": "user", "content": "weather?"}, + {"role": "assistant", "content": [signed, tu]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01A", "content": "sunny"}]}, + { + "role": "assistant", + "content": [signed, {"type": "text", "text": ""}, {"type": "text", "text": "Sunny."}], + }, + ] + out = strip_empty_content_blocks_from_anthropic_messages(msgs) + assert len(out) == 4 + assert out[1] is msgs[1] + assert [b["type"] for b in out[3]["content"]] == ["thinking", "text"] + assert out[3]["content"][0] is signed + assert out[3]["content"][1]["text"] == "Sunny." + def test_strip_thinking_blocks_from_anthropic_messages(self): from litellm.llms.anthropic.common_utils import ( strip_thinking_blocks_from_anthropic_messages, From 4f538f48c84bf147ea5a26b5256961a730b73caf Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 15 Oct 2026 11:49:53 -0700 Subject: [PATCH 04/13] fix(anthropic): keep signed empty thinking blocks and honor thinking.block_binding (#46342) * fix(anthropic): treat adaptive thinking as enabled when mapping response_format is_thinking_enabled only recognised thinking.type=enabled, so a request that paired thinking.type=adaptive with response_format took the synthetic JSON tool path with a forced tool_choice, which Anthropic rejects alongside thinking. The same check also gates the Bedrock Converse and Databricks forced tool_choice downgrades, so they now see adaptive thinking too. * fix(anthropic): keep signed thinking blocks with empty text on chat completions Models that default to display: omitted return signed thinking blocks whose text is empty, and Anthropic rejects the next turn when they are not replayed unchanged. is_unsignable_thinking_block dropped them along with unsigned blocks, and the inline content path required non-empty text, so a tool-use turn on Opus 4.7+ or the 5 families lost its thinking blocks and 400ed on the tool-result turn. Only a missing signature now drops a block, and the mid-conversation system placement mirrors that rule. The /v1/messages strip helper is left to #41362. * fix(anthropic): skip the thinking-strip retry when block_binding asks for an error The /v1/messages retry loop strips thinking blocks and resends on any invalid-thinking 400, including the binding mismatch a caller opted into with thinking.block_binding.prefix_mismatch_behavior: error. The retry decision now sees the request body and surfaces that 400 instead of overriding the caller's choice. #42888 covers the litellm_params side. * fix(bedrock): carry thinking.block_binding through the Invoke thinking rebuilds Bedrock Invoke rebuilt the thinking param from scratch when display: updates was requested, on both the /v1/messages and chat completions paths, and the legacy enabled-to-adaptive translation plus the clear_thinking injection did the same, so thinking.block_binding never reached Bedrock. Every rebuild now spreads the caller's thinking and only overrides the keys it owns, and AnthropicThinkingParam types block_binding. The thinking-binding-controls beta was already on the Bedrock allowlist. * fix(anthropic): keep forcing the response_format tool under adaptive thinking This reverts ad9e861ae7. Anthropic, Bedrock Invoke and Vertex all accept a forced tool_choice next to thinking.type=adaptive and only reject it with type=enabled, so counting adaptive as enabled fixed no live failure. It did drop the forced JSON tool, so Claude answered response_format requests in prose, and it downgraded a required tool_choice to auto on Bedrock Converse * test(e2e): drop the adaptive thinking tool round trip case With adaptive thinking the model often skips thinking on the weather prompt, so turn 1 carries no signed empty block and the case fails before the replay it is meant to cover. Anthropic also answers the replay with a 200 whether or not the block is forwarded, so the case could not tell the fix from main. The factory unit tests pin the outgoing shape instead * fix(anthropic): only skip the thinking-strip retry for a binding mismatch 400 prefix_mismatch_behavior: error asks Anthropic to reject a block bound to another conversation. It says nothing about unrelated bad blocks such as a null signature, so those still get the strip-and-retry self-heal. Also trims the empty-thinking docstrings to what the code guarantees and types the new test parameters * refactor(anthropic): narrow thinking fields with the mapping validator instead of typing.cast typing.cast is banned by the strict ruff gate (TID251), so the three new casts pushed that rule one over its merge-base count. The existing pydantic mapping validator gives the same typed view without an unchecked cast, and it now has a generic name since it is no longer only used for Claude Code payloads * refactor(anthropic): read thinking settings through the validated mapping so basedpyright stays at its merge-base count The Bedrock clear_thinking rewrite still narrowed with a cast and an untyped dict fallback, which added two reportUnknownVariableType errors. It now uses the same pydantic mapping validator, made public for that reuse, and the retry hook validates the request body itself so the handler passes a typed dict instead of an untyped one --------- Co-authored-by: yuneng-jiang (cherry picked from commit f470e418eaa73ef5fc9d001a8a06723364b55cdb) --- .../prompt_templates/common_utils.py | 11 +-- .../prompt_templates/factory.py | 7 +- .../mid_conversation_system.py | 3 +- litellm/llms/anthropic/common_utils.py | 42 ++++++--- .../anthropic_messages/transformation.py | 18 +++- .../anthropic_claude3_transformation.py | 3 +- .../anthropic_claude3_transformation.py | 20 +--- litellm/llms/custom_httpx/llm_http_handler.py | 4 +- litellm/types/llms/anthropic.py | 5 + ...llm_core_utils_prompt_templates_factory.py | 94 +++++++++---------- ...rompt_templates_mid_conversation_system.py | 2 + .../anthropic/test_anthropic_common_utils.py | 74 +++++++++++++++ ...ations_anthropic_claude3_transformation.py | 57 +++++++++++ .../test_anthropic_claude3_transformation.py | 37 ++++++++ .../custom_httpx/test_llm_http_handler.py | 65 +++++++++++++ 15 files changed, 344 insertions(+), 98 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 0243141eda2..378858afa56 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -2100,10 +2100,8 @@ def is_unsignable_thinking_block(block: object) -> bool: Anthropic verifies the signature cryptographically, so a block with a null, empty, or missing signature (e.g. from an open-source reasoning model) is rejected with a 400, and so is a block whose signature or data carries - another provider's encrypted reasoning. It also rejects a `thinking` block - whose text is empty or whitespace-only ("each thinking block must contain - thinking"), regardless of signature, e.g. when a `thinking_blocks` history - item from a non-Anthropic reasoning provider is replayed through this path. + another provider's encrypted reasoning. A signed block is kept whatever its + text, since ``display: omitted`` returns signed blocks with empty ``thinking``. `redacted_thinking` blocks carry no signature and are always kept. """ if is_encrypted_reasoning_block(block): @@ -2114,10 +2112,7 @@ def is_unsignable_thinking_block(block: object) -> bool: if mapping.get("type") != "thinking": return False signature: Final = mapping.get("signature") - if not (isinstance(signature, str) and len(signature) > 0): - return True - thinking_text: Final = mapping.get("thinking") - return not (isinstance(thinking_text, str) and len(thinking_text.strip()) > 0) + return not (isinstance(signature, str) and len(signature) > 0) def strip_encrypted_reasoning_from_messages( diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 3b0cf0ce238..b9a8bf8c83f 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -2774,13 +2774,8 @@ def anthropic_messages_pt( if not isinstance(m, dict): continue # handle thinking blocks - thinking_block = cast(str, m.get("thinking", "")) text_block = cast(str, m.get("text", "")) - if ( - m.get("type", "") == "thinking" - and len(thinking_block) > 0 - and not is_unsignable_thinking_block(m) - ): # don't pass empty text blocks. anthropic api raises errors. + if m.get("type", "") == "thinking" and not is_unsignable_thinking_block(m): anthropic_message: ChatCompletionThinkingBlock | AnthropicMessagesTextParam = cast( ChatCompletionThinkingBlock, m ) diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py index cb4bbf00c85..2851258ccb0 100644 --- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py +++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py @@ -265,8 +265,7 @@ def _assistant_part_renders(part: object) -> bool: """A text part always renders: the converter pads empty text with a placeholder.""" part_type: Final = message_field(part, "type") if part_type == "thinking": - thinking: Final = message_field(part, "thinking") - return isinstance(thinking, str) and bool(thinking) and _thinking_block_renders(part) + return _thinking_block_renders(part) return part_type in _RENDERED_ASSISTANT_PART_TYPES or ( isinstance(part_type, str) and part_type.endswith("_tool_result") ) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index d037f76ed03..12d77112b08 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -106,7 +106,7 @@ _INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$") _DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$") _DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)") _CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:" -_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) +_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) _CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) @@ -155,9 +155,9 @@ def is_claude_code_user_agent(user_agent: str) -> bool: return user_agent.startswith(_CLAUDE_CODE_USER_AGENT_PREFIXES) -def _validated_claude_code_mapping(value: object) -> dict[object, object] | None: +def validated_object_mapping(value: object) -> dict[object, object] | None: try: - return _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER.validate_python(value) + return _OBJECT_MAPPING_ADAPTER.validate_python(value) except ValidationError: return None @@ -194,7 +194,7 @@ def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None: blocks: Final = _validated_claude_code_list(system) if blocks is None: return None - block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks) + block_mappings: Final = tuple(validated_object_mapping(block) for block in blocks) if any(block is None for block in block_mappings): return None text_values: Final = tuple( @@ -227,7 +227,7 @@ def is_claude_code_one_shot_subagent_request( tools: object, user_agent: str | None, ) -> bool: - only_message: Final = _validated_claude_code_mapping(messages[0]) if len(messages) == 1 else None + only_message: Final = validated_object_mapping(messages[0]) if len(messages) == 1 else None return ( user_agent is not None and is_claude_code_user_agent(user_agent) @@ -926,21 +926,23 @@ class AnthropicModelInfo(BaseLLMModelInfo): return if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider): return - thinking: Final = optional_params.get("thinking") - if not isinstance(thinking, dict) or thinking.get("type") != "enabled": + thinking: Final = validated_object_mapping(optional_params.get("thinking")) + if thinking is None or thinking.get("type") != "enabled": return + budget_tokens: Final = thinking.get("budget_tokens") or 0 effort: Final = AnthropicModelInfo._legacy_budget_to_effort( model=model, - budget_tokens=int(thinking.get("budget_tokens") or 0), + budget_tokens=int(budget_tokens) if isinstance(budget_tokens, (int, float, str)) else 0, custom_llm_provider=custom_llm_provider, ) existing_output_config: Final = optional_params.get("output_config") display: Final = thinking.get("display") - if display in ("summarized", "omitted"): - optional_params["thinking"] = {"type": "adaptive", "display": display} - else: - optional_params["thinking"] = {"type": "adaptive"} + supported_display: Final = {"display": display} if display in ("summarized", "omitted") else {} + carried_thinking_keys: Final = { + key: value for key, value in thinking.items() if key not in ("type", "budget_tokens", "display") + } + optional_params["thinking"] = {**carried_thinking_keys, "type": "adaptive", **supported_display} optional_params["output_config"] = { "effort": effort, **(existing_output_config if isinstance(existing_output_config, dict) else MappingProxyType({})), @@ -1608,6 +1610,20 @@ def is_anthropic_invalid_thinking_block_error(error_text: str) -> bool: return "must contain thinking" in lower +def is_anthropic_thinking_binding_mismatch_error(error_text: str) -> bool: + """messages.N.content.M: Invalid `signature` in `thinking` block. The block is bound to a different conversation.""" + return "bound to a different conversation" in error_text.lower() + + +def anthropic_request_errors_on_thinking_binding_mismatch(request_body: object) -> bool: + """``thinking.block_binding.prefix_mismatch_behavior: error`` is the caller asking + for the 400, so a retry that strips thinking must not override it.""" + body: Final = validated_object_mapping(request_body) + thinking: Final = validated_object_mapping(body.get("thinking")) if body is not None else None + binding: Final = validated_object_mapping(thinking.get("block_binding")) if thinking is not None else None + return binding is not None and binding.get("prefix_mismatch_behavior") == "error" + + def strip_thinking_blocks_from_anthropic_messages(messages: Sequence[object]) -> list[object]: """ Return a new message list with thinking / redacted_thinking content blocks removed @@ -1890,7 +1906,7 @@ def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None: def tool_call_is_rebuilt_as_server_tool_use(tool_call_id: object, provider_specific_fields: object) -> bool: - fields: Final = _validated_claude_code_mapping(provider_specific_fields) + fields: Final = validated_object_mapping(provider_specific_fields) if not isinstance(tool_call_id, str) or fields is None: return False return ( diff --git a/litellm/llms/base_llm/anthropic_messages/transformation.py b/litellm/llms/base_llm/anthropic_messages/transformation.py index 5758d156a61..bd6c53a0453 100644 --- a/litellm/llms/base_llm/anthropic_messages/transformation.py +++ b/litellm/llms/base_llm/anthropic_messages/transformation.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod from collections.abc import AsyncIterator -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final import httpx @@ -179,16 +179,28 @@ class BaseAnthropicMessagesConfig(ABC): """ return 2 - def should_retry_anthropic_messages_on_http_error(self, e: httpx.HTTPStatusError, litellm_params: dict) -> bool: + def should_retry_anthropic_messages_on_http_error( + self, e: httpx.HTTPStatusError, litellm_params: dict, request_body: object = None + ) -> bool: """ When True, async_anthropic_messages_handler will transform the request body and issue one more attempt (bounded by max_retry_on_anthropic_messages_http_error). + A binding-mismatch 400 is not retried when the body's ``thinking.block_binding`` + asked for it, since the retry would strip the very blocks the caller pinned. """ from litellm.llms.anthropic.common_utils import ( + anthropic_request_errors_on_thinking_binding_mismatch, is_anthropic_invalid_thinking_block_error, + is_anthropic_thinking_binding_mismatch_error, ) - return e.response.status_code == 400 and is_anthropic_invalid_thinking_block_error(e.response.text) + if e.response.status_code != 400: + return False + error_text: Final[str] = e.response.text + caller_asked_for_this_error: Final = anthropic_request_errors_on_thinking_binding_mismatch( + request_body + ) and is_anthropic_thinking_binding_mismatch_error(error_text) + return not caller_asked_for_this_error and is_anthropic_invalid_thinking_block_error(error_text) def transform_anthropic_messages_request_on_http_error(self, e: httpx.HTTPStatusError, request_data: dict) -> dict: """ diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 0be087d2dd4..b8bc276de0d 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -123,8 +123,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): and isinstance(translated_thinking, dict) and translated_thinking.get("type") == "adaptive" ): - thinking_with_display: Final[AnthropicThinkingParam] = {"type": "adaptive", "display": "updates"} - optional_params["thinking"] = thinking_with_display + optional_params["thinking"] = {**translated_thinking, "display": "updates"} # The stub model hides the original model from the parent's forced-tool-use backstop response_format_tool_choice: Final = optional_params.get("tool_choice") diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 26a0302ef61..eb6ec66bb09 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -17,7 +17,7 @@ from litellm.llms.anthropic.chat.transformation import ( DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, AnthropicConfig, ) -from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.anthropic.common_utils import AnthropicModelInfo, validated_object_mapping from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -49,7 +49,6 @@ from litellm.types.llms.anthropic import ( ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER, ANTHROPIC_TOOL_SEARCH_BETA_HEADER, - AnthropicThinkingParam, ) from litellm.types.llms.bedrock import BedrockInvokeAnthropicMessagesRequest from litellm.types.llms.openai import AllMessageValues @@ -349,18 +348,10 @@ class AmazonAnthropicClaudeMessagesConfig( if not isinstance(output_config, dict): output_config = {} output_config.setdefault("effort", self._effort_from_thinking_budget(budget_tokens)) - thinking: Final = anthropic_messages_request.get("thinking") - display: Final = thinking.get("display") if isinstance(thinking, dict) else None + thinking: Final = validated_object_mapping(anthropic_messages_request.get("thinking")) + carried: Final = {key: value for key, value in (thinking or {}).items() if key not in ("type", "budget_tokens")} anthropic_messages_request["output_config"] = output_config - if display is None: - adaptive_thinking: Final[AnthropicThinkingParam] = {"type": "adaptive"} - anthropic_messages_request["thinking"] = adaptive_thinking - else: - adaptive_thinking_with_display: Final[AnthropicThinkingParam] = { - "type": "adaptive", - "display": display, - } - anthropic_messages_request["thinking"] = adaptive_thinking_with_display + anthropic_messages_request["thinking"] = {**carried, "type": "adaptive"} verbose_logger.debug( "Bedrock clear_thinking_20251015: injected adaptive thinking with effort=%s for model=%s", output_config.get("effort"), @@ -701,8 +692,7 @@ class AmazonAnthropicClaudeMessagesConfig( and isinstance(translated_thinking, dict) and translated_thinking.get("type") == "adaptive" ): - thinking_with_display: Final[AnthropicThinkingParam] = {"type": "adaptive", "display": "updates"} - anthropic_messages_request["thinking"] = thinking_with_display + anthropic_messages_request["thinking"] = {**translated_thinking, "display": "updates"} self._normalize_system_role_messages(anthropic_messages_request, model=model) ######################################################### ############## BEDROCK Invoke SPECIFIC TRANSFORMATION ### diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d87b6c0c854..9cfc47761dd 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2091,7 +2091,7 @@ class BaseLLMHTTPHandler: # str when the caller passes a pre-serialized (unsigned) body to avoid # re-dumping; bytes when a provider signed the request (e.g. Bedrock). signed_json_body: str | bytes | None, - request_body: dict, + request_body: dict[str, object], stream: bool, logging_obj: LiteLLMLoggingObj, provider_config: BaseAnthropicMessagesConfig, @@ -2118,7 +2118,7 @@ class BaseLLMHTTPHandler: except httpx.HTTPStatusError as e: hit_max_attempt = attempt_idx + 1 == max_attempts should_retry = provider_config.should_retry_anthropic_messages_on_http_error( - e=e, litellm_params=litellm_params_dict + e=e, litellm_params=litellm_params_dict, request_body=request_body ) if should_retry and not hit_max_attempt: if logging_obj.baseline_cache_context is not None: diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index f0b5a1bb84e..1e526f48099 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -732,10 +732,15 @@ ANTHROPIC_API_ONLY_HEADERS: Final = { # fails if calling anthropic on vertex ai } +class AnthropicThinkingBlockBinding(TypedDict, total=False): + prefix_mismatch_behavior: ReadOnly[Literal["error", "drop_block"]] + + class AnthropicThinkingParam(TypedDict, total=False): type: ReadOnly[Literal["enabled", "adaptive", "disabled"]] budget_tokens: int display: ReadOnly[Literal["summarized", "omitted", "updates"]] + block_binding: ReadOnly[AnthropicThinkingBlockBinding] class ANTHROPIC_HOSTED_TOOLS(str, Enum): diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 56e97173eb3..5b823af47e9 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -32,7 +32,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( sanitize_messages_for_tool_calling, THOUGHT_SIGNATURE_SEPARATOR, ) -from litellm.types.llms.openai import ChatCompletionToolMessage +from litellm.types.llms.openai import ChatCompletionAssistantMessage, ChatCompletionToolMessage from litellm.utils import ( _invalidate_model_cost_lowercase_map, function_setup, @@ -4130,44 +4130,42 @@ def test_anthropic_messages_pt_drops_a_system_message_with_no_text(): assert [m["role"] for m in result] == ["user", "assistant"] -def test_anthropic_messages_pt_drops_empty_but_signed_thinking_block(): - """ - Anthropic rejects a `thinking` block whose `thinking` text is empty, even - when it carries a valid-looking signature, with: - 400 messages.N.content.M.thinking: each thinking block must contain thinking - This shape is reachable via cross-provider replay of a `thinking_blocks` - history item (see PR #36033), so `is_unsignable_thinking_block()` must - also check the thinking text, not just the signature. - """ +@pytest.mark.parametrize( + "assistant_message", + ( + { + "role": "assistant", + "content": "4", + "thinking_blocks": [{"type": "thinking", "thinking": "", "signature": "sig_abc123_looks_valid"}], + }, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "", "signature": "sig_abc123_looks_valid"}, + {"type": "text", "text": "4"}, + ], + }, + ), + ids=("thinking_blocks", "inline_content"), +) +def test_anthropic_messages_pt_keeps_empty_but_signed_thinking_block( + assistant_message: ChatCompletionAssistantMessage, +) -> None: + """``display: omitted`` returns signed thinking blocks with empty text, and they + must reach the next turn unchanged.""" from litellm.litellm_core_utils.prompt_templates.factory import ( anthropic_messages_pt, ) - messages = [ - {"role": "user", "content": "What's 2+2?"}, - { - "role": "assistant", - "content": "4", - "thinking_blocks": [ - { - "type": "thinking", - "thinking": "", - "signature": "sig_abc123_looks_valid", - } - ], - }, - ] - result = anthropic_messages_pt( - messages=messages, - model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "What's 2+2?"}, assistant_message], + model="claude-opus-5-5", llm_provider="anthropic", ) assistant_msg = result[1] assert isinstance(assistant_msg["content"], list) - content_types = [block.get("type") for block in assistant_msg["content"]] - assert "thinking" not in content_types, "empty-text thinking block must be dropped even though it has a signature" + assert assistant_msg["content"][0] == {"type": "thinking", "thinking": "", "signature": "sig_abc123_looks_valid"} def test_anthropic_messages_pt_keeps_non_empty_signed_thinking_block(): @@ -4281,29 +4279,31 @@ def test_anthropic_messages_pt_drops_unsigned_thinking_block(): assert "thinking" not in content_types, "unsigned thinking block must still be dropped" -def test_is_unsignable_thinking_block_treats_whitespace_only_as_empty(): +def test_anthropic_messages_pt_drops_unsigned_empty_thinking_block(): """ - Edge case: a `thinking` field that is present but whitespace-only (e.g. - a single trailing newline forwarded from another provider's empty - reasoning summary) is functionally empty and Anthropic's API will still - reject it with "each thinking block must contain thinking". We treat it - the same as a fully empty string and drop the block. - - The check lives in the shared `is_unsignable_thinking_block` helper, which - `_drop_unsignable_thinking_blocks` calls standalone, so the whitespace-aware - test has to hold there rather than only at the factory call site. + An empty block without a signature is the shape Claude Code replays after a + non-Anthropic turn, and Anthropic rejects it, so only the missing signature + decides whether a block is dropped. """ - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - is_unsignable_thinking_block, + from litellm.litellm_core_utils.prompt_templates.factory import ( + anthropic_messages_pt, ) - whitespace_only_block = { - "type": "thinking", - "thinking": " \n\t ", - "signature": "sig_abc123_looks_valid", - } + result = anthropic_messages_pt( + messages=[ + {"role": "user", "content": "What's 2+2?"}, + { + "role": "assistant", + "content": [{"type": "thinking", "thinking": "", "signature": ""}, {"type": "text", "text": "4"}], + }, + ], + model="claude-opus-5-5", + llm_provider="anthropic", + ) - assert is_unsignable_thinking_block(whitespace_only_block) is True + assistant_msg = result[1] + assert isinstance(assistant_msg["content"], list) + assert [block.get("type") for block in assistant_msg["content"]] == ["text"] _CONTENT_LESS_USER_MESSAGES: Final = ({"role": "user"}, {"role": "user", "content": None}) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py index 9e3704a9df0..822c0058533 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py @@ -337,6 +337,7 @@ def test_flagged_placement_keeps_a_system_whose_empty_assistant_follower_ends_th }, {"role": "assistant", "content": None, "thinking_blocks": [{"type": "redacted_thinking", "data": "x"}]}, {"role": "assistant", "content": [{"type": "thinking", "thinking": "hm", "signature": "s"}]}, + {"role": "assistant", "content": [{"type": "thinking", "thinking": "", "signature": "s"}]}, {"role": "assistant", "content": None, "function_call": {"name": "f", "arguments": "{}"}}, litellm.Message( role="assistant", @@ -349,6 +350,7 @@ def test_flagged_placement_keeps_a_system_whose_empty_assistant_follower_ends_th "signed-thinking-block", "redacted-thinking-block", "signed-thinking-part", + "signed-empty-thinking-part", "function-call", "pydantic-tool-calls", ], diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index ce6f9754c74..dcf0d58d675 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -14,6 +14,7 @@ import json import os import sys import threading +from collections.abc import Mapping from pathlib import Path from types import SimpleNamespace from typing import Final @@ -1332,6 +1333,12 @@ class TestPassthroughAuthToken: assert url == "https://custom.example.com/v1/messages" +_BOUND_THINKING_ERROR: Final = ( + "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation." +) +_NULL_SIGNATURE_ERROR: Final = "messages.1.content.0.thinking.signature.str: Input should be a valid string" + + class TestAnthropicThinkingSignatureSelfHeal: """Helpers for retrying after invalid thinking blocks in replayed history: invalid encrypted signatures, and blocks with empty thinking text.""" @@ -1997,6 +2004,49 @@ class TestAnthropicThinkingSignatureSelfHeal: assert "thinking" not in data assert data["messages"] == [] + @pytest.mark.parametrize( + ("thinking", "error_message", "expected"), + ( + ( + {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}, + _BOUND_THINKING_ERROR, + False, + ), + ( + {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}, + _NULL_SIGNATURE_ERROR, + True, + ), + ( + {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "drop_block"}}, + _BOUND_THINKING_ERROR, + True, + ), + ({"type": "adaptive"}, _BOUND_THINKING_ERROR, True), + ), + ids=("error", "error_on_unrelated_bad_block", "drop_block", "unbound"), + ) + def test_anthropic_messages_retry_honors_thinking_block_binding( + self, thinking: Mapping[str, object], error_message: str, expected: bool + ) -> None: + from litellm.llms.anthropic.pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + req: Final = httpx.Request("POST", "https://api.anthropic.com/v1/messages") + resp: Final = httpx.Response( + 400, + request=req, + text=json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": error_message}}), + ) + err: Final = httpx.HTTPStatusError("bad", request=req, response=resp) + + should_retry = AnthropicMessagesConfig().should_retry_anthropic_messages_on_http_error( + err, {}, request_body={"model": "claude-opus-5-5", "thinking": thinking} + ) + + assert should_retry is expected + class TestClaudeOpus48AdaptiveThinking: """Opus 4.8 requires adaptive thinking (``thinking.type='adaptive'`` + @@ -4122,6 +4172,30 @@ def test_shared_legacy_thinking_translation_preserves_supported_display( assert optional_params["thinking"] == expected_thinking +@pytest.mark.usefixtures("local_model_cost_map") +def test_shared_legacy_thinking_translation_carries_block_binding() -> None: + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + optional_params: Final = { + "thinking": { + "type": "enabled", + "budget_tokens": 2048, + "block_binding": {"prefix_mismatch_behavior": "error"}, + }, + } + + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model="claude-opus-5", + optional_params=optional_params, + custom_llm_provider="azure_ai", + ) + + assert optional_params["thinking"] == { + "type": "adaptive", + "block_binding": {"prefix_mismatch_behavior": "error"}, + } + + @pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") @pytest.mark.parametrize("action", (None, "tool_addition", "tool_removal")) @pytest.mark.parametrize("explicit_beta", (False, True)) diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 859b57c7cf3..1abebaaa8c5 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -760,6 +760,30 @@ def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_th assert result["output_config"] == {"effort": "high"} +def test_bedrock_chat_invoke_adaptive_thinking_keeps_the_forced_response_format_tool(local_model_cost_map): + """Anthropic accepts a forced tool_choice next to adaptive thinking, and dropping + it lets the model answer in prose instead of the requested JSON.""" + result = AmazonAnthropicClaudeConfig().map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}}, + }, + }, + "thinking": {"type": "adaptive"}, + "max_tokens": 8192, + }, + optional_params={}, + model="us.anthropic.claude-opus-5", + drop_params=False, + ) + + assert result["tool_choice"] == {"type": "tool", "name": result["tools"][0]["name"]} + assert result["thinking"] == {"type": "adaptive"} + + def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking(local_model_cost_map): """Regression: the tool-based ``response_format`` path swaps in a Claude 3 stub model before the shared Anthropic mapping, which hid the adaptive-only model @@ -974,6 +998,39 @@ def test_bedrock_chat_invoke_preserves_display_when_translating_legacy_thinking( assert ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER in result.get("anthropic_beta", []) +BLOCK_BINDING: Final = {"prefix_mismatch_behavior": "error"} + + +@pytest.mark.parametrize( + "requested_thinking", + ( + {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING}, + {"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING}, + ), + ids=("adaptive", "legacy"), +) +def test_bedrock_chat_invoke_keeps_block_binding_with_display_updates( + local_model_cost_map, local_beta_headers_config, requested_thinking +) -> None: + config: Final = AmazonAnthropicClaudeConfig() + model: Final = "us.anthropic.claude-opus-5" + optional_params: Final = config.map_openai_params( + non_default_params={"max_tokens": 512, "thinking": requested_thinking}, + optional_params={}, + model=model, + drop_params=False, + ) + result: Final = config.transform_request( + model=model, + messages=[{"role": "user", "content": "Reply with OK"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result.get("thinking") == {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING} + + FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" EAGER_TOOL_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 52821f58482..579972f8258 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -3529,6 +3529,43 @@ def test_bedrock_clear_thinking_preserves_display_updates() -> None: assert ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER in result.get("anthropic_beta", []) +BLOCK_BINDING: Final = {"prefix_mismatch_behavior": "error"} + + +@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") +@pytest.mark.parametrize( + ("model", "requested_thinking", "extra_params"), + ( + ("eu.anthropic.claude-opus-5", {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING}, {}), + ( + "eu.anthropic.claude-opus-5", + {"type": "enabled", "budget_tokens": 24000, "display": "updates", "block_binding": BLOCK_BINDING}, + {}, + ), + ( + "us.anthropic.claude-opus-4-6", + {"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING}, + {"context_management": {"edits": [{"type": "clear_thinking_20251015"}]}}, + ), + ), + ids=("adaptive", "legacy", "legacy_clear_thinking"), +) +def test_bedrock_messages_keeps_block_binding_with_display_updates( + model: str, requested_thinking: dict[str, object], extra_params: dict[str, object] +) -> None: + from litellm.types.router import GenericLiteLLMParams + + result: Final = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": "Reply with OK"}], + anthropic_messages_optional_request_params={"max_tokens": 512, "thinking": requested_thinking, **extra_params}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("thinking") == {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING} + + @pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") @pytest.mark.parametrize("action", (None, "tool_addition", "tool_removal")) @pytest.mark.parametrize("explicit_beta", (False, True)) diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 25d2d11b5ca..15e1ae6b140 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -1804,6 +1804,71 @@ async def test_anthropic_post_retry_reserializes_mutated_body(): assert "MUTATED" in second_sent # ... the mutated body +_BOUND_THINKING_ERROR: Final = ( + '{"type":"error","error":{"type":"invalid_request_error","message":"messages.1.content.0: ' + 'Invalid `signature` in `thinking` block. The block is bound to a different conversation."}}' +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("prefix_mismatch_behavior", "expected_posts"), + (("error", 1), ("drop_block", 2)), +) +async def test_anthropic_post_skips_the_thinking_strip_retry_when_block_binding_errors( + prefix_mismatch_behavior: str, expected_posts: int +) -> None: + """A caller that pinned ``prefix_mismatch_behavior: error`` asked for the 400, + so the handler must surface it instead of stripping the bound blocks and retrying.""" + from litellm.llms.anthropic.pass_through.messages.transformation import AnthropicMessagesConfig + + handler = BaseLLMHTTPHandler() + thinking_block = {"type": "thinking", "thinking": "", "signature": "sig"} + request_body = { + "model": "claude-opus-5-5", + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [thinking_block, {"type": "text", "text": "hello"}]}, + ], + "thinking": {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": prefix_mismatch_behavior}}, + } + request = httpx.Request("POST", "https://api.anthropic.com/v1/messages") + http_client = Mock() + http_client.post = AsyncMock( + side_effect=[ + httpx.Response(400, request=request, text=_BOUND_THINKING_ERROR), + httpx.Response(200, request=request, json={"type": "message", "content": []}), + ] + ) + logging_obj: Final = Mock(baseline_cache_context=None) + logging_obj.model_call_details = {} + + async def _post() -> httpx.Response: + return await handler._async_post_anthropic_messages_with_http_error_retry( + async_httpx_client=http_client, + request_url="https://api.anthropic.com/v1/messages", + headers={}, + signed_json_body=None, + request_body=request_body, + stream=False, + logging_obj=logging_obj, + provider_config=AnthropicMessagesConfig(), + litellm_params=GenericLiteLLMParams(), + api_key="k", + model="claude-opus-5-5", + ) + + if expected_posts == 1: + with pytest.raises(BaseLLMException) as raised: + await _post() + assert raised.value.status_code == 400 + assert "bound to a different conversation" in raised.value.message + assert request_body["messages"][1]["content"][0] == thinking_block + else: + assert (await _post()).status_code == 200 + assert http_client.post.await_count == expected_posts + + def test_base_responses_config_sign_request_is_noop_by_default(): """Default responses sign_request must be a no-op: unchanged headers, no signed body. From da5b2b556b21789ea8e46f1b58163a7a5d360a4e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 12:51:47 -0700 Subject: [PATCH 05/13] fix(vector-stores): keep proxy context out of provider vector store metadata (#47021) * fix(vector-stores): keep proxy context out of provider vector store metadata The proxy merged its request context (user_api_key_* fields, endpoint, user_agent, requester_ip_address, litellm_api_version) and the router's deployment info into the vector store request's metadata. For vector stores that field is also OpenAI's own metadata, which is capped at 16 keys, so a store created with a team key was full on creation and every later metadata update returned 400. Vector store routes now keep that context in litellm_metadata, like batches, files and responses, and the router does the same for avector_store_create. Provider-side metadata holds only what the caller sent. Adds db_scripts/strip_vector_store_internal_metadata.py, an opt-in script that removes the proxy-written keys from stores created before this change. It is a dry run unless --apply is passed. * fix(vector-stores): page the metadata cleanup script without recursion and tolerate non-string values Listing every store recursed once per page, which overflows the interpreter stack on large projects. Pages are now walked iteratively. Metadata values are no longer required to be strings, and a malformed response exits 2 with a fixed message instead of a traceback that could echo metadata values. * fix(vector-stores): read proxy agent id, tags and spend_logs_metadata from litellm_metadata Vector store routes now keep proxy context in litellm_metadata, but the v3 limiter only read agent_id from metadata, and logging let provider metadata tags or spend_logs_metadata replace the proxy's values. A string tags value also crashed get_request_tags. The limiter now prefers the proxy-resolved agent_id in litellm_metadata. Logging combines list tags from both buckets (metadata first, deduplicated) and merges dict spend_logs_metadata with litellm_metadata winning, ignoring values of the wrong type. * fix(vector-stores): narrow the logging fallback to mistyped caller values The tag union and spend_logs_metadata merge changed logged tags and spend_logs_metadata on every route that has both metadata buckets, deduped tags quadratically, and stopped background interaction polling from keeping string tags as sent. Caller values now win exactly as before, and the proxy's litellm_metadata tags or spend_logs_metadata only replace a caller value of the wrong type. The limiter reads agent ids through a pydantic model instead of a cast and ignores non-string ids. * fix(vector-stores): log the proxy's full tag list on litellm_metadata routes On routes that keep proxy context in litellm_metadata, the proxy builds litellm_metadata.tags from the caller's tags plus key, team and header tags, but get_request_tags preferred the caller's metadata.tags. A caller sending metadata.tags on a vector store call dropped the key's tags from request_tags, which tag spend tracking reads. get_request_tags now prefers a non-empty litellm_metadata tag list, and falls back to metadata.tags as before. * fix(vector-stores): strip LiteLLM-written keys at the provider boundary instead of moving metadata buckets Moving vector store proxy context into litellm_metadata broke shared readers (agent limits, request tags, merged logging metadata), and fixing them widened the change to every litellm_metadata route. This reverts the bucket move, the router change and the reader fixes, so proxy context stays in metadata exactly as on main. Vector store create and update now drop the keys LiteLLM itself writes (user_api_key*, requester and deployment keys, agent ids) from the metadata sent to the provider, in the shared optional-param helper both paths use. The cleanup script reuses the same key list. * fix(vector-stores): pass non-string metadata keys through the provider metadata filter * fix(vector-stores): strip LiteLLM-written metadata keys only on proxy and router calls * fix(vector-stores): send the provider exactly the metadata the caller sent through the proxy * fix(vector-stores): snapshot caller metadata from the raw request body and keep it out of logged request bodies * fix(vector-stores): copy the parsed-body metadata when the raw request body cannot be reparsed (cherry picked from commit 4846cfabc7f4b8cf6a83a960b7d1123eebef4315) --- .github/workflows/test-unit.yml | 1 + .../strip_vector_store_internal_metadata.py | 201 ++++++++++++++++++ litellm/constants.py | 1 + litellm/proxy/litellm_pre_call_utils.py | 3 +- .../proxy/vector_store_endpoints/endpoints.py | 9 + litellm/vector_stores/utils.py | 91 +++++++- tests/unit/db_scripts/__init__.py | 0 ...st_strip_vector_store_internal_metadata.py | 198 +++++++++++++++++ .../test_vector_store_endpoints.py | 157 ++++++++++++++ .../unit/test_router/test_router_endpoints.py | 37 ++++ tests/unit/vector_stores/test_main.py | 90 +++++++- 11 files changed, 782 insertions(+), 6 deletions(-) create mode 100644 db_scripts/strip_vector_store_internal_metadata.py create mode 100644 tests/unit/db_scripts/__init__.py create mode 100644 tests/unit/db_scripts/test_strip_vector_store_internal_metadata.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index ae8f8cb240f..3f1b6b11d35 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -670,6 +670,7 @@ jobs: - shard: proxy-extras test-path: >- tests/unit/litellm_proxy_extras + tests/unit/db_scripts python-version: "3.12" workers: 2 reruns: 0 diff --git a/db_scripts/strip_vector_store_internal_metadata.py b/db_scripts/strip_vector_store_internal_metadata.py new file mode 100644 index 00000000000..34ffa85dcd1 --- /dev/null +++ b/db_scripts/strip_vector_store_internal_metadata.py @@ -0,0 +1,201 @@ +import argparse +import os +import sys +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from itertools import accumulate, repeat, takewhile +from typing import Final, Literal, TextIO + +import httpx +from pydantic import BaseModel, ConfigDict, JsonValue, ValidationError + +from litellm.vector_stores.utils import is_litellm_written_metadata_key + +DEFAULT_API_BASE: Final = "https://api.openai.com/v1" +LIST_PAGE_SIZE: Final = "100" + + +class VectorStore(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + metadata: dict[str, JsonValue] | None = None + + +class VectorStorePage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + data: tuple[VectorStore, ...] + has_more: bool = False + + +class OpenAIErrorBody(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + message: str = "" + + +class OpenAIErrorResponse(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + error: OpenAIErrorBody + + +@dataclass(frozen=True, slots=True) +class Options: + api_key: str + api_base: str + vector_store_ids: tuple[str, ...] + extra_keys: frozenset[str] + apply: bool + + +@dataclass(frozen=True, slots=True) +class StoreResult: + store_id: str + matched_keys: tuple[str, ...] + status: Literal["dry_run", "updated", "failed"] + error: str = "" + + +def is_proxy_written_key(key: str, extra_keys: frozenset[str]) -> bool: + return is_litellm_written_metadata_key(key) or key in extra_keys + + +def matching_keys(store: VectorStore, extra_keys: frozenset[str]) -> tuple[str, ...]: + return tuple(sorted(key for key in (store.metadata or {}) if is_proxy_written_key(key, extra_keys))) + + +def fetch_page(client: httpx.Client, after: str | None) -> VectorStorePage: + params: Final = {"limit": LIST_PAGE_SIZE} if after is None else {"limit": LIST_PAGE_SIZE, "after": after} + return VectorStorePage.model_validate_json(client.get("/vector_stores", params=params).raise_for_status().content) + + +def fetch_next_page(client: httpx.Client, page: VectorStorePage | None) -> VectorStorePage | None: + if page is None or not page.has_more or not page.data: + return None + return fetch_page(client, after=page.data[-1].id) + + +def list_all_stores(client: httpx.Client) -> Iterator[VectorStore]: + pages: Final = accumulate( + repeat(None), lambda page, _: fetch_next_page(client, page), initial=fetch_page(client, after=None) + ) + for page in takewhile(lambda page: page is not None, pages): + yield from page.data + + +def retrieve_stores(client: httpx.Client, store_ids: tuple[str, ...]) -> Iterator[VectorStore]: + for store_id in store_ids: + yield VectorStore.model_validate_json(client.get(f"/vector_stores/{store_id}").raise_for_status().content) + + +def describe_http_error(response: httpx.Response) -> str: + try: + message: Final = OpenAIErrorResponse.model_validate_json(response.content).error.message + except ValidationError: + return f"HTTP {response.status_code}" + return f"HTTP {response.status_code}: {message}" + + +def delete_keys(client: httpx.Client, store_id: str, keys: tuple[str, ...]) -> StoreResult: + try: + response: Final = client.post(f"/vector_stores/{store_id}", json={"metadata": dict.fromkeys(keys)}) + except httpx.TransportError as error: + return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=type(error).__name__) + if response.is_success: + return StoreResult(store_id=store_id, matched_keys=keys, status="updated") + return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=describe_http_error(response)) + + +def clean_store(client: httpx.Client, store: VectorStore, options: Options) -> StoreResult | None: + keys: Final = matching_keys(store, options.extra_keys) + if not keys: + return None + if not options.apply: + return StoreResult(store_id=store.id, matched_keys=keys, status="dry_run") + return delete_keys(client, store.id, keys) + + +def format_result(result: StoreResult) -> str: + key_list: Final = ", ".join(result.matched_keys) + count: Final = len(result.matched_keys) + match result.status: + case "dry_run": + return f"{result.store_id}: would delete {count} keys: {key_list}" + case "updated": + return f"{result.store_id}: deleted {count} keys: {key_list}" + case "failed": + return f"{result.store_id}: failed to delete {count} keys ({result.error}): {key_list}" + + +def clean_and_report( + client: httpx.Client, stores: Iterator[VectorStore], options: Options, out: TextIO +) -> Iterator[StoreResult | None]: + for store in stores: + result: Final = clean_store(client, store, options) + if result is not None: + print(format_result(result), file=out) + yield result + + +def run(options: Options, client: httpx.Client, out: TextIO) -> int: + stores: Final = ( + retrieve_stores(client, options.vector_store_ids) if options.vector_store_ids else list_all_stores(client) + ) + scanned: Final = tuple(clean_and_report(client, stores, options, out)) + results: Final = tuple(result for result in scanned if result is not None) + updated: Final = sum(1 for result in results if result.status == "updated") + failed: Final = sum(1 for result in results if result.status == "failed") + mode: Final = "applied" if options.apply else "dry run, pass --apply to delete" + print( + f"scanned={len(scanned)} matched={len(results)} updated={updated} failed={failed} ({mode})", + file=out, + ) + return 1 if failed else 0 + + +def parse_options(argv: Sequence[str] | None) -> Options: + parser: Final = argparse.ArgumentParser( + description="Remove LiteLLM proxy-written keys from OpenAI vector store metadata. Dry run by default." + ) + parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY"), help="defaults to $OPENAI_API_KEY") + parser.add_argument("--api-base", default=DEFAULT_API_BASE) + parser.add_argument("--vector-store-id", action="append", default=[], help="repeatable, defaults to every store") + parser.add_argument("--extra-key", action="append", default=[], help="repeatable, another key name to delete") + parser.add_argument("--apply", action="store_true", help="delete the matching keys instead of listing them") + args: Final = parser.parse_args(argv) + if not args.api_key: + parser.error("--api-key or OPENAI_API_KEY is required") + return Options( + api_key=args.api_key, + api_base=args.api_base.rstrip("/"), + vector_store_ids=tuple(args.vector_store_id), + extra_keys=frozenset(args.extra_key), + apply=args.apply, + ) + + +def main(argv: Sequence[str] | None = None, transport: httpx.BaseTransport | None = None) -> int: + options: Final = parse_options(argv) + with httpx.Client( + base_url=options.api_base, + headers={"Authorization": f"Bearer {options.api_key}"}, + timeout=60.0, + transport=transport, + ) as client: + try: + return run(options, client, sys.stdout) + except httpx.HTTPStatusError as error: + sys.stderr.write(f"could not read vector stores: {describe_http_error(error.response)}\n") + return 2 + except httpx.TransportError as error: + sys.stderr.write(f"could not reach {options.api_base}: {type(error).__name__}\n") + return 2 + except ValidationError: + sys.stderr.write(f"unexpected vector store response shape from {options.api_base}\n") + return 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/litellm/constants.py b/litellm/constants.py index 3132cedcc8b..6f812ac5388 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1588,6 +1588,7 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags" +VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata" ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4e75226f1e4..5b9835d667e 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -31,6 +31,7 @@ from litellm.constants import ( SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, TYPESAFE_DANGER_ZONE_HEADER_PREFIX, + VECTOR_STORE_CALLER_METADATA_KWARG, X_LITELLM_DISABLE_CALLBACKS, ) from litellm.integrations.custom_guardrail import without_server_streaming_classification @@ -2099,7 +2100,7 @@ def refresh_proxy_server_request_body_snapshot( if not isinstance(proxy_server_request, dict): return _body_snapshot_exclude: Final = ( - frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj"}) + frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj", VECTOR_STORE_CALLER_METADATA_KWARG}) | _TRANSPORT_ONLY_CREDENTIAL_KEYS | _CALLBACK_CREDENTIAL_KEYS ) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 22721dc26bc..1b54dddcf4f 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -9,12 +9,14 @@ from typing import ( from fastapi import APIRouter, Depends, HTTPException, Request, Response +from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.http_parsing_utils import read_raw_json_body from litellm.proxy.utils import jsonify_object from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, @@ -23,6 +25,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( ) from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse +from litellm.vector_stores.utils import VectorStoreWriteBody, callers_metadata_snapshot from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry router: Final = APIRouter() @@ -213,6 +216,9 @@ async def vector_store_create( ) data: Final = await read_request_body(request=request) + data[VECTOR_STORE_CALLER_METADATA_KWARG] = callers_metadata_snapshot( + await read_raw_json_body(request), VectorStoreWriteBody.model_validate(data) + ) # Check for target_model_names parameter target_model_names: Final = data.pop("target_model_names", None) @@ -448,6 +454,9 @@ async def vector_store_update( data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line if "vector_store_id" not in data: data["vector_store_id"] = vector_store_id + data[VECTOR_STORE_CALLER_METADATA_KWARG] = callers_metadata_snapshot( + await read_raw_json_body(request), VectorStoreWriteBody.model_validate(data) + ) data = await _update_request_data_with_litellm_managed_vector_store_registry( data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict diff --git a/litellm/vector_stores/utils.py b/litellm/vector_stores/utils.py index c0ac4edbf74..cf12f28642f 100644 --- a/litellm/vector_stores/utils.py +++ b/litellm/vector_stores/utils.py @@ -1,11 +1,92 @@ +import copy +from collections.abc import Container, Mapping from typing import Any, Final, cast, get_type_hints +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from litellm.constants import ROUTING_REQUEST_TAGS_METADATA_KEY, VECTOR_STORE_CALLER_METADATA_KWARG from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.types.vector_stores import ( VectorStoreCreateOptionalRequestParams, VectorStoreSearchOptionalRequestParams, ) +LITELLM_WRITTEN_METADATA_KEYS: Final = frozenset( + { + "requester_ip_address", + "user_agent", + "endpoint", + "litellm_api_version", + "deployment", + "deployment_model_name", + "api_base", + "agent_id", + "actor_agent_id", + "target_agent_id", + "billing_agent_id", + "agent_execution_mode", + "verified_human_user_id", + "litellm_caller_session_id", + } +) + + +def is_litellm_written_metadata_key(key: str) -> bool: + return key == "user_api_key" or key.startswith("user_api_key_") or key in LITELLM_WRITTEN_METADATA_KEYS + + +def _is_callers_key(key: object) -> bool: + return not isinstance(key, str) or not is_litellm_written_metadata_key(key) + + +def _sent_by_proxy_or_router(params: Mapping[str, object], metadata: Container[object]) -> bool: + return params.get("proxy_server_request") is not None or ROUTING_REQUEST_TAGS_METADATA_KEY in metadata + + +def _callers_metadata(metadata: Mapping[object, object]) -> Mapping[object, object]: + return {key: value for key, value in metadata.items() if _is_callers_key(key)} + + +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) + + +class VectorStoreWriteBody(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + metadata: object = None + + +def _json_object(raw: str | bytes) -> Mapping[str, object] | None: + try: + return _JSON_OBJECT.validate_json(raw) + except ValidationError: + return None + + +def _decoded_metadata(metadata: object) -> object: + if not isinstance(metadata, str): + return metadata + decoded: Final = _json_object(metadata) + return metadata if decoded is None else decoded + + +def callers_metadata_snapshot(raw_json_body: bytes | None, parsed_body: VectorStoreWriteBody) -> object: + raw_body: Final = None if raw_json_body is None else _json_object(raw_json_body) + metadata: Final = copy.deepcopy(parsed_body.metadata) if raw_body is None else raw_body.get("metadata") + return _decoded_metadata(metadata) + + +def _metadata_for_provider(params: Mapping[str, object]) -> object: + if VECTOR_STORE_CALLER_METADATA_KWARG in params: + return params[VECTOR_STORE_CALLER_METADATA_KWARG] + metadata: Final = params.get("metadata") + if not isinstance(metadata, Mapping): + return metadata + request_metadata: Final = cast(Mapping[object, object], metadata) # cast-ok: request metadata is a JSON object + if not _sent_by_proxy_or_router(params, request_metadata): + return request_metadata + return _callers_metadata(request_metadata) + class VectorStoreRequestUtils: """Helper utils for constructing Vector Store search requests""" @@ -48,7 +129,11 @@ class VectorStoreRequestUtils: Returns: VectorStoreCreateOptionalRequestParams instance with only the valid parameters """ - valid_keys: Final = get_type_hints(VectorStoreCreateOptionalRequestParams).keys() + valid_keys: Final = get_type_hints(VectorStoreCreateOptionalRequestParams).keys() - {"metadata"} filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None} - - return cast(VectorStoreCreateOptionalRequestParams, filtered_params) + metadata: Final = _metadata_for_provider(params) + if metadata is None: + return cast(VectorStoreCreateOptionalRequestParams, filtered_params) + return cast( # cast-ok: same params as above plus the metadata the provider should store + VectorStoreCreateOptionalRequestParams, {**filtered_params, "metadata": metadata} + ) diff --git a/tests/unit/db_scripts/__init__.py b/tests/unit/db_scripts/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/db_scripts/test_strip_vector_store_internal_metadata.py b/tests/unit/db_scripts/test_strip_vector_store_internal_metadata.py new file mode 100644 index 00000000000..05b055f75a1 --- /dev/null +++ b/tests/unit/db_scripts/test_strip_vector_store_internal_metadata.py @@ -0,0 +1,198 @@ +import importlib.util +import json +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import httpx +import pytest + +_SCRIPT: Final = Path(__file__).resolve().parents[3] / "db_scripts" / "strip_vector_store_internal_metadata.py" +_spec: Final = importlib.util.spec_from_file_location("strip_vector_store_internal_metadata", _SCRIPT) +strip_script: Final = importlib.util.module_from_spec(_spec) +sys.modules[_spec.name] = strip_script +_spec.loader.exec_module(strip_script) + +POLLUTED: Final = { + "id": "vs_1", + "metadata": { + "user_api_key": "hash-value-1", + "user_api_key_hash": "hash-value-1", + "user_api_key_user_email": "member@example.com", + "endpoint": "http://proxy/v1/vector_stores", + "deployment": "openai/gpt-5", + "team": "llmproxy", + "trace_id": "trace-123", + }, +} +CLEAN: Final = {"id": "vs_2", "metadata": {"team": "llmproxy"}} +NO_METADATA: Final = {"id": "vs_3", "metadata": None} + + +def _fake_openai( + stores: tuple[Mapping[str, object], ...], page_size: int = 100, post_status: int = 200 +) -> tuple[httpx.MockTransport, list[httpx.Request]]: + seen: Final[list[httpx.Request]] = [] + by_id: Final = {store["id"]: store for store in stores} + + def handle(request: httpx.Request) -> httpx.Response: + seen.append(request) + path: Final = request.url.path.removeprefix("/v1") + if request.method == "POST": + if post_status != 200: + return httpx.Response(post_status, json={"error": {"message": "upstream rejected the update"}}) + return httpx.Response(200, json=by_id[path.removeprefix("/vector_stores/")]) + if path == "/vector_stores": + after: Final = request.url.params.get("after") + start: Final = 0 if after is None else [store["id"] for store in stores].index(after) + 1 + page: Final = stores[start : start + page_size] + return httpx.Response(200, json={"data": list(page), "has_more": start + page_size < len(stores)}) + return httpx.Response(200, json=by_id[path.removeprefix("/vector_stores/")]) + + return httpx.MockTransport(handle), seen + + +def _run(args: list[str], transport: httpx.MockTransport) -> int: + return strip_script.main(["--api-key", "sk-test", *args], transport=transport) + + +def _posts(seen: list[httpx.Request]) -> list[tuple[str, dict[str, object]]]: + return [(request.url.path, json.loads(request.content)) for request in seen if request.method == "POST"] + + +def test_dry_run_lists_matching_key_names_without_values_and_sends_no_update(capsys): + transport, seen = _fake_openai((POLLUTED, CLEAN)) + + exit_code: Final = _run([], transport) + + out: Final = capsys.readouterr().out + assert exit_code == 0 + assert ( + "vs_1: would delete 5 keys: deployment, endpoint, user_api_key, user_api_key_hash, user_api_key_user_email" + in out + ) + assert "scanned=2 matched=1 updated=0 failed=0" in out + assert "hash-value-1" not in out + assert "member@example.com" not in out + assert _posts(seen) == [] + assert {request.headers["authorization"] for request in seen} == {"Bearer sk-test"} + + +def test_apply_nulls_exactly_the_proxy_written_keys_and_leaves_caller_keys_alone(capsys): + transport, seen = _fake_openai((POLLUTED, CLEAN, NO_METADATA)) + + exit_code: Final = _run(["--apply"], transport) + + assert exit_code == 0 + assert _posts(seen) == [ + ( + "/v1/vector_stores/vs_1", + { + "metadata": { + "deployment": None, + "endpoint": None, + "user_api_key": None, + "user_api_key_hash": None, + "user_api_key_user_email": None, + } + }, + ) + ] + assert "scanned=3 matched=1 updated=1 failed=0" in capsys.readouterr().out + + +def test_extra_key_is_deleted_alongside_the_built_in_keys(): + transport, seen = _fake_openai((POLLUTED,)) + + _run(["--apply", "--extra-key", "trace_id"], transport) + + assert _posts(seen)[0][1]["metadata"]["trace_id"] is None + assert "team" not in _posts(seen)[0][1]["metadata"] + + +def test_listing_follows_has_more_with_the_last_store_id(capsys): + transport, seen = _fake_openai((CLEAN, POLLUTED, NO_METADATA), page_size=2) + + _run([], transport) + + assert [request.url.params.get("after") for request in seen] == [None, "vs_1"] + assert "scanned=3 matched=1" in capsys.readouterr().out + + +def test_listing_walks_more_pages_than_the_interpreter_recursion_limit(capsys): + stores: Final = tuple({"id": f"vs_{index}", "metadata": None} for index in range(sys.getrecursionlimit() + 50)) + transport, _ = _fake_openai(stores, page_size=1) + + assert _run([], transport) == 0 + assert f"scanned={len(stores)} matched=0" in capsys.readouterr().out + + +def test_non_string_metadata_values_are_scanned_without_printing_them(capsys): + transport, _ = _fake_openai(({"id": "vs_9", "metadata": {"user_api_key": "secret-value", "count": 7}},)) + + exit_code: Final = _run([], transport) + + out: Final = capsys.readouterr().out + assert exit_code == 0 + assert "vs_9: would delete 1 keys: user_api_key" in out + assert "secret-value" not in out + + +def test_an_unexpected_response_shape_exits_non_zero_without_echoing_metadata(capsys): + def malformed(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": [{"metadata": {"user_api_key": "secret-value"}}], "has_more": False}) + + exit_code: Final = _run([], httpx.MockTransport(malformed)) + + captured: Final = capsys.readouterr() + assert exit_code == 2 + assert "unexpected vector store response shape" in captured.err + assert "secret-value" not in captured.out + captured.err + + +def test_vector_store_id_retrieves_only_the_named_stores_without_listing(capsys): + transport, seen = _fake_openai((POLLUTED, CLEAN, NO_METADATA)) + + _run(["--vector-store-id", "vs_2", "--vector-store-id", "vs_1"], transport) + + assert [request.url.path for request in seen] == ["/v1/vector_stores/vs_2", "/v1/vector_stores/vs_1"] + assert "scanned=2 matched=1" in capsys.readouterr().out + + +def test_a_failed_update_is_reported_and_exits_non_zero(capsys): + transport, _ = _fake_openai((POLLUTED,), post_status=400) + + exit_code: Final = _run(["--apply"], transport) + + out: Final = capsys.readouterr().out + assert exit_code == 1 + assert "vs_1: failed to delete 5 keys (HTTP 400: upstream rejected the update)" in out + assert "updated=0 failed=1" in out + + +def test_a_listing_error_exits_non_zero_without_updating_anything(): + seen: Final[list[httpx.Request]] = [] + + def reject(request: httpx.Request) -> httpx.Response: + seen.append(request) + return httpx.Response(401, json={"error": {"message": "bad key"}}) + + assert _run(["--apply"], httpx.MockTransport(reject)) == 2 + assert [request.method for request in seen] == ["GET"] + + +def test_an_unreachable_api_exits_non_zero(): + def unreachable(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("connection refused", request=request) + + assert _run([], httpx.MockTransport(unreachable)) == 2 + + +def test_an_api_key_is_required(monkeypatch): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + + with pytest.raises(SystemExit) as exit_info: + strip_script.main([]) + + assert exit_info.value.code == 2 diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 2f70c1f5188..48642a9b50d 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -6,6 +6,7 @@ import pytest from fastapi import HTTPException, Request import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) @@ -3256,3 +3257,159 @@ class TestConfigOwnedVectorStores: assert response["status"] == "success", response prisma.db.litellm_managedvectorstorestable.delete.assert_awaited_once_with(where={"vector_store_id": self.DB_ID}) assert registry.get_litellm_managed_vector_store_from_registry(self.DB_ID) is None + + +CALLERS_METADATA = {"team": "llmproxy", "endpoint": "ingest", "trace_id": "callers-trace"} + + +class _PreCallBodies(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged_bodies = [] + + def log_pre_api_call(self, model, messages, kwargs): + self.logged_bodies.append(kwargs["litellm_params"]["proxy_server_request"]["body"]) + + +@pytest.mark.parametrize("operation", ["create", "update"]) +@pytest.mark.parametrize("model", [None, "gpt-5"]) +@pytest.mark.parametrize( + ("metadata_field", "expected_metadata"), + [ + ({"metadata": CALLERS_METADATA}, CALLERS_METADATA), + ({"metadata": {}}, {}), + ({"metadata": json.dumps(CALLERS_METADATA)}, CALLERS_METADATA), + ({"metadata": None}, None), + ({}, None), + ], +) +def test_proxy_vector_store_writes_send_the_provider_exactly_the_callers_metadata( + monkeypatch, operation, model, metadata_field, expected_metadata +): + import httpx + import respx + from fastapi.testclient import TestClient + + from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import app + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setattr(proxy_server, "master_key", "sk-master-test") + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[{"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5", "api_key": "sk-test"}}] + ), + ) + pre_call_bodies = _PreCallBodies() + monkeypatch.setattr(litellm, "callbacks", [pre_call_bodies]) + sent_bodies = [] + + def echo(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + sent_bodies.append(body) + return httpx.Response( + 200, + json={ + "id": "vs_1", + "object": "vector_store", + "created_at": 0, + "name": "x", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": body.get("metadata"), + }, + ) + + path = "/v1/vector_stores" if operation == "create" else "/v1/vector_stores/vs_1" + routing = {} if model is None else {"model": model} + with respx.mock(assert_all_called=False) as upstream: + upstream.post(url__regex=r"https://api\.openai\.com/v1/vector_stores.*").mock(side_effect=echo) + response = TestClient(app).post( + path, + json={"name": "x", **routing, **metadata_field}, + headers={ + "Authorization": "Bearer sk-master-test", + "x-litellm-tags": "header-tag", + "x-litellm-trace-id": "header-trace", + "x-litellm-session-id": "header-session", + }, + ) + + assert response.status_code == 200, response.text + assert len(sent_bodies) == 1 + assert sent_bodies[0].get("metadata") == expected_metadata + assert len(pre_call_bodies.logged_bodies) == 1 + assert "name" in pre_call_bodies.logged_bodies[0] + assert VECTOR_STORE_CALLER_METADATA_KWARG not in pre_call_bodies.logged_bodies[0] + + +@pytest.mark.parametrize("operation", ["create", "update"]) +@pytest.mark.parametrize( + "request_body", + [ + {"data": {"name": "x", "metadata": json.dumps(CALLERS_METADATA)}}, + { + "content": json.dumps({"name": "x", "metadata": CALLERS_METADATA})[:-1] + + ', "nested": ' + + "[" * 300 + + "]" * 300 + + "}", + "headers": {"Content-Type": "application/json"}, + }, + ], + ids=["form", "json-too-deep-to-reparse"], +) +def test_proxy_vector_store_writes_send_only_the_callers_metadata_when_the_raw_body_cannot_be_reparsed( + monkeypatch, operation, request_body +): + import httpx + import respx + from fastapi.testclient import TestClient + + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import app + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setattr(proxy_server, "master_key", "sk-master-test") + monkeypatch.setattr(proxy_server, "llm_router", litellm.Router(model_list=[])) + sent_bodies = [] + + def echo(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + sent_bodies.append(body) + return httpx.Response( + 200, + json={ + "id": "vs_1", + "object": "vector_store", + "created_at": 0, + "name": "x", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": body.get("metadata"), + }, + ) + + path = "/v1/vector_stores" if operation == "create" else "/v1/vector_stores/vs_1" + headers = { + "Authorization": "Bearer sk-master-test", + "x-litellm-trace-id": "header-trace", + **request_body.get("headers", {}), + } + body_kwargs = {key: value for key, value in request_body.items() if key != "headers"} + with respx.mock(assert_all_called=False) as upstream: + upstream.post(url__regex=r"https://api\.openai\.com/v1/vector_stores.*").mock(side_effect=echo) + response = TestClient(app).post(path, headers=headers, **body_kwargs) + + assert response.status_code == 200, response.text + assert len(sent_bodies) == 1 + assert sent_bodies[0].get("metadata") == CALLERS_METADATA diff --git a/tests/unit/test_router/test_router_endpoints.py b/tests/unit/test_router/test_router_endpoints.py index 74a1bfefeee..815b32b33f7 100644 --- a/tests/unit/test_router/test_router_endpoints.py +++ b/tests/unit/test_router/test_router_endpoints.py @@ -1,10 +1,14 @@ import asyncio +import json import os +from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch +import httpx import litellm import pytest from litellm import Router +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @pytest.fixture @@ -573,6 +577,39 @@ async def test_init_vector_store_api_endpoints(): mock_original_function.assert_called_once_with(vector_store_id="test-store", custom_llm_provider="openai") +@pytest.mark.asyncio +async def test_avector_store_create_sends_only_the_callers_metadata_to_the_provider(): + def echo_vector_store(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "vs_echo", + "object": "vector_store", + "created_at": 0, + "name": body["name"], + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": body.get("metadata"), + }, + ) + + router: Final = Router( + model_list=[{"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5", "api_key": "sk-test"}}] + ) + + response: Final = await router.avector_store_create( + model="gpt-5", + name="x", + metadata={"team": "llmproxy"}, + client=AsyncHTTPHandler(transport=httpx.MockTransport(echo_vector_store)), + ) + + assert response["id"] == "vs_echo" + assert response["metadata"] == {"team": "llmproxy"} + + def test_apply_default_settings(): """ Test the apply_default_settings method. diff --git a/tests/unit/vector_stores/test_main.py b/tests/unit/vector_stores/test_main.py index 1c968126c42..25a378270b1 100644 --- a/tests/unit/vector_stores/test_main.py +++ b/tests/unit/vector_stores/test_main.py @@ -8,16 +8,19 @@ model_dump() it (the #19550 serialization trap). """ import json +from collections.abc import Mapping +from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest import litellm.vector_stores.main as vector_stores_main from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, ) -from litellm.llms.custom_httpx.http_handler import HTTPHandler -from litellm.vector_stores.main import search +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.vector_stores.main import acreate, aupdate, search MOCK_SEARCH_RESPONSE = { "object": "vector_store.search_results.page", @@ -114,3 +117,86 @@ def test_search_forwards_top_level_user_context_to_bedrock_retrieve(): posted = json.loads(client.post.call_args.kwargs["data"]) assert posted["userContext"] == {"userId": "alice@example.com"} assert posted["retrievalQuery"] == {"text": "q"} + + +PROXY_STAMPED_METADATA: Final = { + "team": "llmproxy", + "user_api_key": "hashed-key", + "user_api_key_hash": "hashed-key", + "user_api_key_team_id": "team-1", + "requester_ip_address": "127.0.0.1", + "user_agent": "curl/8.7.1", + "endpoint": "http://localhost:4000/v1/vector_stores", + "litellm_api_version": "1.0.0", + "deployment": "openai/gpt-5", + "deployment_model_name": "gpt-5", + "api_base": "https://api.openai.com/v1", + "agent_id": "agent-1", + "actor_agent_id": "agent-2", + "target_agent_id": "agent-3", + "billing_agent_id": "agent-4", + "agent_execution_mode": "autonomous", + "verified_human_user_id": "user-1", + "litellm_caller_session_id": "session-1", + "user_api_keyboard": "qwerty", + 1: "a", +} + + +def _copy(metadata: Mapping[object, object] | None) -> dict[object, object] | None: + return None if metadata is None else dict(metadata) + + +async def _metadata_sent_to_provider( + operation: str, metadata: Mapping[object, object] | None, **kwargs: object +) -> list[object]: + sent_metadata: Final[list[object]] = [] + + def echo_vector_store(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + sent_metadata.append(body.get("metadata")) + return httpx.Response( + 200, + json={ + "id": "vs_1", + "object": "vector_store", + "created_at": 0, + "name": "x", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": body.get("metadata"), + }, + ) + + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(echo_vector_store)) + if operation == "create": + await acreate(name="x", metadata=_copy(metadata), api_key="sk-test", client=client, **kwargs) + else: + await aupdate(vector_store_id="vs_1", metadata=_copy(metadata), api_key="sk-test", client=client, **kwargs) + return sent_metadata + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["create", "update"]) +async def test_proxy_vector_store_writes_send_only_the_callers_metadata_to_the_provider(operation: str): + sent_metadata: Final = await _metadata_sent_to_provider( + operation, + PROXY_STAMPED_METADATA, + proxy_server_request={"url": "http://localhost:4000/v1/vector_stores", "method": "POST", "body": {}}, + ) + + assert sent_metadata == [{"team": "llmproxy", "user_api_keyboard": "qwerty", "1": "a"}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["create", "update"]) +@pytest.mark.parametrize( + "callers_metadata", [{"team": "llmproxy", "endpoint": "ingest", "deployment": "prod", "agent_id": "bot-7"}, None] +) +async def test_direct_sdk_vector_store_writes_send_the_callers_metadata_unchanged( + operation: str, callers_metadata: Mapping[object, object] | None +): + sent_metadata: Final = await _metadata_sent_to_provider(operation, callers_metadata) + + assert sent_metadata == [callers_metadata] From 07f72a67b1d97038e159fe0f26b8d8c52294754b Mon Sep 17 00:00:00 2001 From: Mrunmayee Daware <165297479+Mrun25@users.noreply.github.com> Date: Sun, 18 Oct 2026 12:48:11 +0530 Subject: [PATCH 06/13] fix(proxy): authorize the public name of a team-scoped model on batch retrieve/cancel (#47067) * fix(proxy): authorize the public name of a team-scoped model on batch retrieve/cancel A unified batch id resolves to the deployment's internal model_name (model_name_{team_id}_{uuid}), which is in no key or team allowlist, so authorize_model_for_key always denied and team keys got 403 retrieving or cancelling their own batches. Authorize against model_info team_public_model_name when the resolved deployment is team-scoped, falling back to the resolved name otherwise. Fixes #46958 * fix(proxy): authorize team-scoped batch models by their public name on both auth paths * fix(proxy): accept the owning team's public name for its own team deployment on every auth path Managed batch, file, fine-tuning, vector store, character and video ids resolve to the team deployment's internal name again, so teams that granted the internal name keep access. The model access check now also accepts the owning team's public name and its legacy team aliases for that internal name, and only when the caller's team owns the deployment. * fix(proxy): match a team deployment's public and alias names against the caller's own grants --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> (cherry picked from commit e979bd0ecf836976373e12624585962dd3356f3a) --- litellm/proxy/auth/auth_checks.py | 30 +++++ tests/unit/proxy/auth/test_auth_utils.py | 118 +++++++++++++++++ .../proxy/batches_endpoints/test_endpoints.py | 125 +++++++++++++++++- 3 files changed, 272 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1435033e1be..89841f667db 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4597,6 +4597,24 @@ def check_model_access_helper( _check_model_access_helper: Final = check_model_access_helper +def _owning_team_names_for_team_deployment( + model: str, + llm_router: Router | None, + team_id: str | None, + team_model_aliases: Mapping[str, str] | None, +) -> tuple[str, ...]: + if llm_router is None or team_id is None or not model.startswith(f"model_name_{team_id}_"): + return () + deployment: Final = llm_router.get_deployment_by_model_group_name(model_group_name=model) + public_name: Final = ( + deployment.model_info.team_public_model_name + if deployment is not None and deployment.model_info.team_id == team_id + else None + ) + legacy_aliases: Final = tuple(alias for alias, target in (team_model_aliases or {}).items() if target == model) + return (*((public_name,) if public_name else ()), *legacy_aliases) + + def _can_object_call_model( model: str | list[str], llm_router: Router | None, @@ -4668,6 +4686,13 @@ def _can_object_call_model( *((global_or_router_alias_target,) if global_or_router_alias_target else ()), ) ) + owning_team_names: Final = ( + () + if key_alias_applied + else _owning_team_names_for_team_deployment( + model=model, llm_router=llm_router, team_id=team_id, team_model_aliases=team_model_aliases + ) + ) ## check model access for alias + underlying model - allow if either is in allowed models for m in potential_models: @@ -4679,6 +4704,11 @@ def _can_object_call_model( team_id=team_id, ): return True + if any( + check_model_access_helper(model=name, llm_router=llm_router, models=models, team_id=team_id) + for name in owning_team_names + ): + return True internal_message: Final = ( f"{object_type} not allowed to access model. This {object_type} can only access models={models}. " diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 92bd5734353..5820cf7447f 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -1354,6 +1354,124 @@ async def test_managed_batch_routes_pass_team_model_access_check(route, request_ ) +_TEAM_DEPLOYMENT_ID: Final = "team-a-deployment-id" +_TEAM_INTERNAL_MODEL_NAME: Final = "model_name_team-a_6d1f3c2e-0000-4000-8000-000000000000" +_TEAM_PUBLIC_MODEL_NAME: Final = "team-batch-model" +_TEAM_MANAGED_BATCH_ID: Final = _encode_managed_id( + f"litellm_proxy;model_id:{_TEAM_DEPLOYMENT_ID};llm_batch_id:provider-batch-789" +) + + +def _team_scoped_router(public_model_name: str | None = _TEAM_PUBLIC_MODEL_NAME) -> Router: + return Router( + model_list=[ + { + "model_name": _TEAM_INTERNAL_MODEL_NAME, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + "model_info": { + "id": _TEAM_DEPLOYMENT_ID, + "team_id": "team-a", + **({"team_public_model_name": public_model_name} if public_model_name else {}), + }, + }, + ] + ) + + +def _team_scoped_character_id() -> str: + from litellm.types.videos.utils import encode_character_id_with_provider + + return encode_character_id_with_provider(character_id="char_team", provider="openai", model_id=_TEAM_DEPLOYMENT_ID) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route, request_data", + [ + pytest.param("/v1/batches/{batch_id}", {"batch_id": _TEAM_MANAGED_BATCH_ID}, id="batch-retrieve"), + pytest.param("/v1/batches/{batch_id}/cancel", {"batch_id": _TEAM_MANAGED_BATCH_ID}, id="batch-cancel"), + pytest.param( + "/v1/fine_tuning/jobs/{fine_tuning_job_id}", + { + "fine_tuning_job_id": _encode_managed_id( + f"litellm_proxy;model_id:{_TEAM_DEPLOYMENT_ID};generic_response_id:ftjob-1" + ) + }, + id="fine-tuning-job", + ), + pytest.param( + "/v1/vector_stores/{vector_store_id}", + { + "vector_store_id": _encode_managed_id( + "litellm_proxy:vector_store;unified_id,vs-unified;" + f"target_model_names,{_TEAM_PUBLIC_MODEL_NAME};resource_id,vs_1;model_id,{_TEAM_DEPLOYMENT_ID}" + ) + }, + id="vector-store", + ), + pytest.param( + "/v1/videos/characters/{character_id}", + {"character_id": _team_scoped_character_id()}, + id="character", + ), + ], +) +async def test_team_scoped_managed_ids_authorize_only_the_owning_team_by_public_name( + route: str, request_data: dict[str, str] +) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_checks import can_key_call_resolved_model + + llm_router: Final = _team_scoped_router() + + async def check(team_id: str) -> None: + await can_key_call_resolved_model( + model=get_model_from_request( + request_data=request_data, route=route, llm_router=llm_router, team_id=team_id + ), + llm_model_list=None, + valid_token=UserAPIKeyAuth(api_key="sk-team", team_id=team_id, team_models=[_TEAM_PUBLIC_MODEL_NAME]), + llm_router=llm_router, + ) + + await check("team-a") + with pytest.raises(ProxyException) as exc_info: + await check("team-b") + assert exc_info.value.code == "403" + + +@pytest.mark.asyncio +async def test_legacy_team_alias_authorizes_a_team_batch_only_for_keys_granted_the_alias() -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_checks import can_key_call_resolved_model + + llm_router: Final = _team_scoped_router(public_model_name=None) + + async def check(key_models: list[str]) -> None: + await can_key_call_resolved_model( + model=get_model_from_request( + request_data={"batch_id": _TEAM_MANAGED_BATCH_ID}, + route="/v1/batches/{batch_id}", + llm_router=llm_router, + team_id="team-a", + ), + llm_model_list=None, + valid_token=UserAPIKeyAuth( + api_key="sk-team", + team_id="team-a", + models=key_models, + team_models=["legacy-team-model"], + team_model_aliases={"legacy-team-model": _TEAM_INTERNAL_MODEL_NAME}, + ), + llm_router=llm_router, + ) + + await check(["legacy-team-model"]) + with pytest.raises(ProxyException) as exc_info: + await check(["some-other-model"]) + assert exc_info.value.code == "403" + + def test_get_model_from_request_resolves_character_id_model_with_router(): from litellm.types.videos.utils import encode_character_id_with_provider diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index 3469dbee748..1580ec60271 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -35,7 +35,7 @@ import logging from contextlib import ExitStack from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -3259,6 +3259,129 @@ async def test_cancel__unified_batch_id_rejects_key_without_model_grant(cancel_h cancel_harness.router_acancel.assert_not_called() +TEAM_INTERNAL_MODEL_NAME: Final = "model_name_team-a_00000000-0000-0000-0000-000000000000" +TEAM_PUBLIC_MODEL_NAME: Final = "team-gpt4o-mini" + + +def _team_scoped_deployment() -> Deployment: + return Deployment( + model_name=TEAM_INTERNAL_MODEL_NAME, + litellm_params=LiteLLM_Params(model="gpt-4o-mini"), + model_info=ModelInfo(id="dep-team-a-1", team_id="team-a", team_public_model_name=TEAM_PUBLIC_MODEL_NAME), + ) + + +def _route_unified_batch_to_team_deployment( + harness: RetrieveHarness | CancelHarness, deployment: Deployment | None = None +) -> None: + target: Final = deployment or _team_scoped_deployment() + harness.router.resolve_model_name_from_model_id = MagicMock(return_value=target.model_name) + harness.router.get_deployment_by_model_group_name = MagicMock(return_value=target) + + +@pytest.mark.asyncio +async def test_retrieve__team_scoped_deployment_authorizes_owning_team_by_public_name( + retrieve_harness: RetrieveHarness, +) -> None: + _route_unified_batch_to_team_deployment(retrieve_harness) + + resp: Final = await call_retrieve( + retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to(TEAM_PUBLIC_MODEL_NAME) + ) + + assert retrieve_harness.router_aretrieve.call_count == 1 + assert resp._hidden_params["model_id"] == "gpt-4o-mini" + + +@pytest.mark.asyncio +async def test_cancel__team_scoped_deployment_authorizes_owning_team_by_public_name( + cancel_harness: CancelHarness, +) -> None: + _route_unified_batch_to_team_deployment(cancel_harness) + + await call_cancel(cancel_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to(TEAM_PUBLIC_MODEL_NAME)) + + assert cancel_harness.router_acancel.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment, user", + [ + pytest.param( + _team_scoped_deployment(), _key_restricted_to("some-other-model"), id="owning-team-key-without-the-model" + ), + pytest.param( + _team_scoped_deployment(), + UserAPIKeyAuth( + api_key="sk-team-b", + team_id="team-b", + team_models=[TEAM_PUBLIC_MODEL_NAME], + models=[TEAM_PUBLIC_MODEL_NAME], + ), + id="other-team-key-granted-the-same-public-name", + ), + pytest.param( + Deployment( + model_name="model_name_team-a_ml_00000000-0000-0000-0000-000000000000", + litellm_params=LiteLLM_Params(model="gpt-4o-mini"), + model_info=ModelInfo( + id="dep-team-a-ml-1", team_id="team-a_ml", team_public_model_name=TEAM_PUBLIC_MODEL_NAME + ), + ), + _key_restricted_to(TEAM_PUBLIC_MODEL_NAME), + id="team-whose-id-prefixes-the-owning-team-id", + ), + ], +) +async def test_retrieve__team_scoped_deployment_rejects_keys_outside_its_grant( + retrieve_harness: RetrieveHarness, deployment: Deployment, user: UserAPIKeyAuth +) -> None: + _route_unified_batch_to_team_deployment(retrieve_harness, deployment) + + with pytest.raises(ProxyException) as exc_info: + await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=user) + + assert exc_info.value.code == "403" + retrieve_harness.router_aretrieve.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment, user", + [ + pytest.param( + _team_scoped_deployment(), + _key_restricted_to(TEAM_INTERNAL_MODEL_NAME), + id="owning-team-key-granted-the-internal-name", + ), + pytest.param( + Deployment( + model_name=TEAM_INTERNAL_MODEL_NAME, + litellm_params=LiteLLM_Params(model="gpt-4o-mini"), + model_info=ModelInfo(id="dep-team-a-1", team_id="team-a"), + ), + UserAPIKeyAuth( + api_key="sk-legacy", + team_id="team-a", + team_models=["legacy-team-model"], + models=["legacy-team-model"], + team_model_aliases={"legacy-team-model": TEAM_INTERNAL_MODEL_NAME}, + ), + id="legacy-team-model-reached-through-a-team-alias", + ), + ], +) +async def test_retrieve__team_scoped_deployment_authorizes_internal_name_and_legacy_alias_grants( + retrieve_harness: RetrieveHarness, deployment: Deployment, user: UserAPIKeyAuth +) -> None: + _route_unified_batch_to_team_deployment(retrieve_harness, deployment) + + await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=user) + + assert retrieve_harness.router_aretrieve.call_count == 1 + + @pytest.mark.asyncio async def test_cancel__executed_batch_rejects_key_without_model_grant(cancel_harness, executed_runner): runner, factory = executed_runner From 978176398aa747f94d56583b84cface08f3f4036 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 19 Oct 2026 15:42:02 -0700 Subject: [PATCH 07/13] Merge pull request #47389 from BerriAI/litellm_backport_count_tokens_stable_1_104_x fix(count_tokens): backport October 19 fixes to stable/1.104.x (cherry picked from commit 5e4d160423016d171b8a3c5a8ff792e8c1db7c86) --- ci_cd/generate_model_prices_schema.py | 20 + .../gitlab/gitlab_prompt_manager.py | 2 +- litellm/litellm_core_utils/core_helpers.py | 26 +- .../prompt_templates/factory.py | 22 + .../anthropic/count_tokens/token_counter.py | 27 +- .../pass_through/adapters/transformation.py | 2 +- .../pass_through/convertible_blocks.py | 106 + .../anthropic/pass_through/messages/utils.py | 7 + .../responses_adapters/transformation.py | 17 +- .../anthropic/count_tokens/handler.py | 11 +- .../anthropic/count_tokens/token_counter.py | 39 +- .../anthropic/count_tokens/transformation.py | 9 +- litellm/llms/base_llm/base_utils.py | 6 +- .../count_tokens/bedrock_token_counter.py | 4 +- litellm/llms/bedrock/count_tokens/kimi_k3.py | 502 ++++ litellm/llms/gemini/common_utils.py | 45 +- litellm/llms/gemini/count_tokens/handler.py | 44 +- .../openai/responses/count_tokens/handler.py | 8 +- .../responses/count_tokens/token_counter.py | 142 +- .../responses/count_tokens/transformation.py | 88 +- litellm/llms/vertex_ai/common_utils.py | 49 +- .../llms/vertex_ai/count_tokens/context.py | 242 ++ .../llms/vertex_ai/count_tokens/handler.py | 32 +- .../llms/vertex_ai/gemini/transformation.py | 8 +- ...odel_prices_and_context_window_backup.json | 28 +- .../proxy/anthropic_endpoints/endpoints.py | 51 +- litellm/proxy/common_request_processing.py | 21 + .../common_utils/token_count_settings.py | 164 ++ litellm/proxy/google_endpoints/endpoints.py | 172 +- litellm/proxy/proxy_server.py | 495 +++- .../proxy/response_api_endpoints/endpoints.py | 5 +- litellm/types/llms/vertex_ai.py | 6 +- litellm/types/token_count.py | 16 + litellm/types/utils.py | 13 + litellm/utils.py | 3 + model_prices_and_context_window.json | 28 +- model_prices_and_context_window.schema.json | 12 + .../test_count_tokens_error_shapes_wire.py | 215 ++ .../test_openai_count_tokens_tools_wire.py | 393 ++++ .../test_vertex_gemini_count_tokens_wire.py | 596 +++++ tests/integration/proxy_config.yaml | 2 + .../integration/spend/test_cache_and_quota.py | 8 +- .../pass_through/test_convertible_blocks.py | 192 ++ .../llms/anthropic/test_count_tokens_oauth.py | 45 + .../unit/llms/azure_ai/anthropic/__init__.py | 0 .../anthropic/count_tokens/__init__.py | 0 .../count_tokens/test_token_counter.py | 217 ++ .../count_tokens/test_bedrock_kimi_k3.py | 524 +++++ .../llms/gemini/test_gemini_common_utils.py | 238 +- .../count_tokens/test_token_counter.py | 436 ++++ ...test_openai_count_tokens_transformation.py | 109 + .../openai/test_openai_workload_identity.py | 53 + .../llms/vertex_ai/count_tokens/__init__.py | 0 .../vertex_ai/count_tokens/test_context.py | 507 ++++ .../vertex_ai/test_vertex_ai_common_utils.py | 237 +- .../anthropic_endpoints/test_endpoints.py | 15 +- .../common_utils/test_token_count_settings.py | 190 ++ .../test_google_api_endpoints.py | 65 +- .../test_model_management_endpoints.py | 48 + .../proxy/proxy_server/test_routes_utils.py | 8 +- tests/unit/proxy/test_proxy_token_counter.py | 2064 ++++++++++++++--- tests/unit/test_count_tokens_public_api.py | 2 +- tests/unit/test_model_prices_schema.py | 46 + .../unit/test_router_model_cost_isolation.py | 34 + tests/unit/test_utils.py | 7 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 +- 66 files changed, 8081 insertions(+), 654 deletions(-) create mode 100644 litellm/llms/anthropic/pass_through/convertible_blocks.py create mode 100644 litellm/llms/bedrock/count_tokens/kimi_k3.py create mode 100644 litellm/llms/vertex_ai/count_tokens/context.py create mode 100644 litellm/proxy/common_utils/token_count_settings.py create mode 100644 litellm/types/token_count.py create mode 100644 tests/integration/providers/test_count_tokens_error_shapes_wire.py create mode 100644 tests/integration/providers/test_openai_count_tokens_tools_wire.py create mode 100644 tests/integration/providers/test_vertex_gemini_count_tokens_wire.py create mode 100644 tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py create mode 100644 tests/unit/llms/azure_ai/anthropic/__init__.py create mode 100644 tests/unit/llms/azure_ai/anthropic/count_tokens/__init__.py create mode 100644 tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py create mode 100644 tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py create mode 100644 tests/unit/llms/openai/responses/count_tokens/test_token_counter.py create mode 100644 tests/unit/llms/vertex_ai/count_tokens/__init__.py create mode 100644 tests/unit/llms/vertex_ai/count_tokens/test_context.py create mode 100644 tests/unit/proxy/common_utils/test_token_count_settings.py diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 8eec07dadda..a5a4b7f6f3b 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -19,6 +19,7 @@ NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0} NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0} BOOLEAN: JsonSchema = {"type": "boolean"} STRING: JsonSchema = {"type": "string"} +NON_BLANK_STRING: Final[JsonSchema] = {"type": "string", "pattern": r"\S"} TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"} WEEKDAY_PATTERN: Final = ( r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$" @@ -103,6 +104,16 @@ OBJECT_KEYS: dict[str, JsonSchema] = { "type": "object", "description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).", }, + "token_count_model": { + "type": "object", + "required": ["model"], + "additionalProperties": False, + "properties": { + "model": NON_BLANK_STRING, + "region": NON_BLANK_STRING, + }, + "description": "Model that counts tokens for this one, and the AWS region its count API runs in.", + }, } ARRAY_KEYS: dict[str, JsonSchema] = { @@ -189,6 +200,10 @@ INTEGER_KEYS: dict[str, JsonSchema] = { **NONNEG_INTEGER, "description": "Smallest prefix the provider will actually cache; absent means the provider default applies.", }, + "token_count_framing": { + **NONNEG_INTEGER, + "description": "Tokens added once per request to a local count estimate, for the provider's chat framing.", + }, "tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."}, "rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."}, } @@ -244,6 +259,11 @@ def cost_schema(key: str) -> JsonSchema: def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]: return { + "token_count_tokenizer": { + "type": "string", + "description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.", + "enum": ["kimi_k3"], + }, "litellm_provider": { "type": "string", "description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.", diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index 113dea541fb..341c8ae271f 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -219,7 +219,7 @@ class GitLabTemplateManager: elif value.replace(".", "").isdigit(): try: result[key] = float(value) - except Exception: + except ValueError: result[key] = value else: result[key] = value.strip("\"'") diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index e857461785c..2cb892f0162 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -3,9 +3,9 @@ import copy import logging import re -from collections.abc import Collection, Iterable, Mapping +from collections.abc import Collection, Iterable, Iterator, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable import httpx from pydantic import TypeAdapter, ValidationError @@ -52,6 +52,28 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non _DROP_PARAMS_BOOL: Final = TypeAdapter(bool) +def as_str_mapping(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, Mapping) else None + + +@runtime_checkable +class _SupportsIter(Protocol): + def __iter__(self) -> Iterator[object]: ... + + +@runtime_checkable +class _SupportsGetItem(Protocol): + def __getitem__(self, index: int, /) -> object: ... + + +def iter_items(value: object) -> Iterable[object]: + if isinstance(value, _SupportsIter): + return value + if isinstance(value, _SupportsGetItem): + return iter(value) + raise TypeError(f"'{type(value).__name__}' object is not iterable") + + def normalize_drop_params(value: object) -> bool | None: if value is None or isinstance(value, bool): return value diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index b9a8bf8c83f..0c42697b6f9 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5776,3 +5776,25 @@ def resolve_structured_messages( if structured: return [msg if isinstance(msg, dict) else msg.model_dump() for msg in structured] return None + + +def is_flat_responses_tool(tool: Mapping[str, object]) -> bool: + tool_type: Final = tool.get("type") + return tool_type in ("function", "custom") and tool_type not in tool and "input_schema" not in tool + + +def chat_completion_tools_for(tool: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + if is_flat_responses_tool(tool): + chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + [dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in a flat Responses tool shape + ) + return tuple(chat_tools) + anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai( + [dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in an Anthropic or chat tool shape + ) + return tuple(anthropic_tools) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 401c09b580b..e05ef8dd83b 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -4,11 +4,14 @@ Anthropic Token Counter implementation using the CountTokens API. from typing import Any, Final +from pydantic import JsonValue + +import litellm from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler() @@ -30,8 +33,10 @@ class AnthropicTokenCounter(BaseTokenCounter): contents: list[dict[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, + tools: list[dict[str, JsonValue]] | None = None, + system: JsonValue = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using Anthropic's CountTokens API. @@ -57,14 +62,24 @@ class AnthropicTokenCounter(BaseTokenCounter): try: auth_header: Final = await AnthropicModelInfo.aget_auth_header( - api_key=litellm_params.get("api_key"), + api_key=litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key, api_base=api_base, litellm_params=litellm_params, allow_workload_identity=True, ) if auth_header is None: - verbose_logger.warning("No Anthropic credential found for token counting") - return None + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="anthropic_api", + error=True, + error_message=( + f"No Anthropic API key is set for model {request_model or model_to_use}, so tokens can't be " + "counted. Set api_key on the deployment or ANTHROPIC_API_KEY in the proxy environment." + ), + status_code=401, + ) result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request( model=model_to_use, diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index 2ee23852894..132ee876cf8 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -461,7 +461,7 @@ class LiteLLMAnthropicMessagesAdapter: *, custom_llm_provider: str | None = None, preserve_midturn_system: bool = False, - ) -> list: + ) -> list[AllMessageValues]: new_messages: Final[list[AllMessageValues]] = [] replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) leading_count: Final = next( diff --git a/litellm/llms/anthropic/pass_through/convertible_blocks.py b/litellm/llms/anthropic/pass_through/convertible_blocks.py new file mode 100644 index 00000000000..543a6efd4c4 --- /dev/null +++ b/litellm/llms/anthropic/pass_through/convertible_blocks.py @@ -0,0 +1,106 @@ +from collections.abc import Iterator, Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages +from litellm.types.utils import TokenCountMessagesFormat + +_CONVERTIBLE_BLOCK_TYPES: Final = MappingProxyType( + { + "user": frozenset({"text", "image", "document", "tool_result"}), + "assistant": frozenset({"text", "tool_use", "thinking", "redacted_thinking"}), + } +) +_SOURCED_BLOCK_TYPES: Final = frozenset({"image", "document"}) +_TOOL_RESULT_BLOCK_TYPES: Final = frozenset({"text", "image", "document"}) + + +_JSON_MESSAGES: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) + + +def _json_messages(messages: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, JsonValue], ...]: + try: + return _JSON_MESSAGES.validate_python(messages) + except ValidationError: + return () + + +def _source_label(block_type: str, source: JsonValue) -> str | None: + if not isinstance(source, dict): + return f"{block_type} without a source" + match source.get("type"): + case "base64" if source.get("data"): + return None + case "url" if source.get("url"): + return None + case "base64" | "url" as source_type: + return f"{block_type} with an empty {source_type} source" + case str() as source_type: + return f"{block_type} with a {source_type} source" + case _: + return f"{block_type} with an untyped source" + + +def _block_label(block: JsonValue, convertible: frozenset[str]) -> str | None: + if not isinstance(block, dict): + return None + block_type: Final = block.get("type") + if not isinstance(block_type, str): + return "a block without a type" + if block_type not in convertible: + return block_type + return _source_label(block_type, block.get("source")) if block_type in _SOURCED_BLOCK_TYPES else None + + +def _tool_result_labels(block: JsonValue) -> Iterator[str]: + if not isinstance(block, dict) or block.get("type") != "tool_result": + return iter(()) + inner: Final = block.get("content") + if not isinstance(inner, list): + return iter(()) + return ( + f"{label} inside a tool_result" + for item in inner + if (label := _block_label(item, _TOOL_RESULT_BLOCK_TYPES)) is not None + ) + + +def _content_labels(block: JsonValue, convertible: frozenset[str]) -> Iterator[str]: + label: Final = _block_label(block, convertible) + return chain(() if label is None else (label,), _tool_result_labels(block)) + + +def _message_labels(message: Mapping[str, JsonValue]) -> Iterator[str]: + role: Final = message.get("role") + convertible: Final = _CONVERTIBLE_BLOCK_TYPES.get(role) if isinstance(role, str) else None + content: Final = message.get("content") + if convertible is None or not isinstance(content, list): + return iter(()) + return chain.from_iterable(_content_labels(block, convertible) for block in content) + + +def unconvertible_anthropic_blocks(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + normalized: Final = _json_messages(normalize_native_messages(messages)) + return tuple(sorted(set(chain.from_iterable(_message_labels(message) for message in normalized)))) + + +def unconvertible_blocks_error( + messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat +) -> str | None: + match messages_format: + case "openai": + return None + case "anthropic": + labels: Final = unconvertible_anthropic_blocks(messages or ()) + return ( + "count_tokens can't convert these content blocks for this model, so the count would leave them out: " + + ", ".join(labels) + if labels + else None + ) + case _: + assert_never(messages_format) diff --git a/litellm/llms/anthropic/pass_through/messages/utils.py b/litellm/llms/anthropic/pass_through/messages/utils.py index 8371615baea..38791ff887e 100644 --- a/litellm/llms/anthropic/pass_through/messages/utils.py +++ b/litellm/llms/anthropic/pass_through/messages/utils.py @@ -160,6 +160,13 @@ def prepare_native_messages( ) +def normalize_native_messages(messages: Sequence[object]) -> list[dict[str, object]]: + """Apply the block rewrites every /v1/messages request gets before it reaches a provider.""" + return flatten_unencrypted_web_search_results_in_anthropic_messages( + sanitize_tool_use_ids_in_anthropic_messages(strip_empty_content_blocks_from_anthropic_messages(messages)) + ) + + @lru_cache(maxsize=1) def anthropic_messages_optional_param_keys() -> frozenset[str]: """ diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py index 642f1857058..b36c3ef25d5 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py @@ -495,6 +495,19 @@ class LiteLLMAnthropicToResponsesAPIAdapter: result["summary"] = "detailed" return result + @staticmethod + def system_blocks_to_instructions(system: object) -> str: + if not isinstance(system, list): + return "" + return "\n".join( + text + for block in system + if isinstance(block, dict) + and block.get("type") == "text" + and isinstance(text := block.get("text"), str) + and text + ) + def translate_request( self, anthropic_request: AnthropicMessagesRequest, @@ -544,9 +557,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(system, str): responses_kwargs["instructions"] = system elif isinstance(system, list): - responses_kwargs["instructions"] = "\n".join( - filter(None, (b.get("text", "") for b in system if isinstance(b, dict) and b.get("type") == "text")) - ) + responses_kwargs["instructions"] = self.system_blocks_to_instructions(system) # max_tokens -> max_output_tokens max_tokens: Final = anthropic_request.get("max_tokens") diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 3270fb3534a..22e467b0de3 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -29,7 +29,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): self, model: str, messages: list[dict[str, Any]], - api_key: str, + api_key: str | None, api_base: str, litellm_params: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, @@ -79,6 +79,15 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): api_key=api_key, litellm_params=litellm_params, ) + if not {"x-api-key", "Authorization"} & headers.keys(): + raise AnthropicError( + status_code=401, + message=( + f"No Azure AI API key or Entra ID credential is set for model {model}, so tokens can't be " + "counted. Set api_key or an Entra ID credential on the deployment, or AZURE_AI_API_KEY or " + "AZURE_API_KEY in the proxy environment." + ), + ) # Use LiteLLM's async httpx client async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.AZURE_AI) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index c31d2c427bb..5de5240ee47 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -5,17 +5,35 @@ Azure AI Anthropic Token Counter implementation using the CountTokens API. import os from typing import Any, Final +from pydantic import JsonValue + +import litellm from litellm._logging import verbose_logger from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler() +def _missing_api_base_response(request_model: str, model_to_use: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="azure_ai_anthropic_api", + error=True, + error_message=( + f"No Azure AI API base is set for model {request_model or model_to_use}, so tokens can't be counted. " + "Set api_base on the deployment or AZURE_AI_API_BASE or AZURE_API_BASE in the proxy environment." + ), + status_code=500, + ) + + class AzureAIAnthropicTokenCounter(BaseTokenCounter): """Token counter implementation for Azure AI Anthropic provider using the CountTokens API.""" @@ -34,6 +52,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using Azure AI Anthropic's CountTokens API. @@ -61,18 +81,15 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): if not api_key: api_key = os.getenv("AZURE_AI_API_KEY") - # Get API base from deployment config or environment - api_base = litellm_params.get("api_base") - if not api_base: - api_base = os.getenv("AZURE_AI_API_BASE") - - if not api_key: - verbose_logger.warning("No Azure AI API key found for token counting") - return None + api_base: Final = ( + litellm_params.get("api_base") + or os.getenv("AZURE_AI_API_BASE") + or os.getenv("AZURE_API_BASE") + or litellm.api_base + ) if not api_base: - verbose_logger.warning("No Azure AI API base found for token counting") - return None + return _missing_api_base_response(request_model, model_to_use) try: result: Final = await azure_ai_anthropic_count_tokens_handler.handle_count_tokens_request( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index d7455439151..5946a1f4480 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -24,7 +24,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): def get_required_headers( self, - api_key: str, + api_key: str | None, litellm_params: dict[str, object] | None = None, ) -> dict[str, str]: """ @@ -45,12 +45,12 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): "Content-Type": "application/json", "anthropic-version": "2023-06-01", "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, - "x-api-key": api_key, # Azure AI Anthropic requires this header + **({"x-api-key": api_key} if api_key else {}), } # Also set up Azure auth headers for flexibility litellm_params = litellm_params or {} - if "api_key" not in litellm_params: + if api_key and "api_key" not in litellm_params: litellm_params["api_key"] = api_key litellm_params_obj: Final = GenericLiteLLMParams.model_validate(litellm_params) @@ -58,8 +58,9 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): # Get Azure auth headers (api-key or Authorization) azure_headers = BaseAzureLLM.base_validate_azure_environment(headers={}, litellm_params=litellm_params_obj) - # Merge Azure auth headers headers.update(azure_headers) + if "api-key" in azure_headers: + headers.setdefault("x-api-key", azure_headers["api-key"]) return headers diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index ba4dad40d13..64a10f777a0 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -9,11 +9,11 @@ from collections.abc import Mapping from typing import Any, Final from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # SDK parser internals -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue from litellm._logging import verbose_logger from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk -from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountResponse +from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountMessagesFormat, TokenCountResponse class BaseTokenCounter(ABC): @@ -27,6 +27,8 @@ class BaseTokenCounter(ABC): request_model: str = "", tools: list[dict[str, Any]] | None = None, system: Any | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: pass diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index ed61707f99e..65a6aee027c 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -16,7 +16,7 @@ from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountResponse +from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" @@ -106,6 +106,8 @@ class BedrockTokenCounter(BaseTokenCounter): request_model: str = "", tools: Sequence[Mapping[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: if not messages: return None diff --git a/litellm/llms/bedrock/count_tokens/kimi_k3.py b/litellm/llms/bedrock/count_tokens/kimi_k3.py new file mode 100644 index 00000000000..5e3fd56e4ba --- /dev/null +++ b/litellm/llms/bedrock/count_tokens/kimi_k3.py @@ -0,0 +1,502 @@ +""" +Exact Bedrock token counts for Kimi K3, which has no CountTokens support. + +Renders the Converse body litellm sends through Moonshot's Kimi K3 chat template and counts it with the +official tokenizer, downloaded once from Hugging Face at a pinned revision and checked against its SHA-256. +""" + +import asyncio +import base64 +import copy +import hashlib +import json +import threading +import time +from collections.abc import Callable, Iterable, Mapping, Sequence +from concurrent.futures import Future +from dataclasses import dataclass +from enum import Enum +from itertools import chain, groupby +from pathlib import Path +from typing import Annotated, Final, Literal, TypeAlias + +import regex +import tiktoken +from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError, model_validator +from typing_extensions import Self + +import litellm +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.prompt_templates.factory import ( + chat_completion_tools_for, + is_flat_responses_tool, +) +from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam + +KIMI_K3_REPO: Final = "moonshotai/Kimi-K3" +KIMI_K3_REVISION: Final = "f831ab66814297da540d832a5235f8e904f29d06" +KIMI_K3_RANKS_FILE: Final = "tiktoken.model" +KIMI_K3_RANKS_SHA256: Final = "b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103" +_PATTERN: Final = "|".join( + ( + r"[\p{Han}]+", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?", + r"\p{N}{1,3}", + r" ?[^\s\p{L}\p{N}]+[\r\n]*", + r"\s*[\r\n]+", + r"\s+(?!\S)", + r"\s+", + ) +) +# Ends chunks only where the pattern above always splits, so their counts sum to the whole text's count while BPE +# stays linear. Only a run of over 1024 characters with no such split, like 2000 spaces, is cut mid-piece. +_CHUNK: Final = regex.compile( + r"(?s).{1,1024}(?:(?<=\S)(?= \S)|(?<=\p{N})(?=\P{N})|(?<=[^\s\p{N}])(?=\p{N})|(?<=\p{L})(?=[^\p{L}\p{M}'])" + r"|(?<=[\r\n])(?=\S)|\Z)|.{1024}\p{N}*" +) + + +@dataclass(frozen=True, slots=True) +class KimiK3TokenizerSource: + ranks_path: Callable[[], str] + sha256: str + + +def _hugging_face_ranks_path() -> str: + import huggingface_hub + + return huggingface_hub.hf_hub_download( + repo_id=KIMI_K3_REPO, + filename=KIMI_K3_RANKS_FILE, + revision=KIMI_K3_REVISION, + local_files_only=litellm.disable_hf_tokenizer_download is True, + ) + + +HUGGING_FACE_SOURCE: Final = KimiK3TokenizerSource(ranks_path=_hugging_face_ranks_path, sha256=KIMI_K3_RANKS_SHA256) + + +def _loaded(source: KimiK3TokenizerSource) -> tiktoken.Encoding | str: + try: + data: Final = Path(source.ranks_path()).read_bytes() + except Exception as e: # noqa: BLE001 # a download or read failure of any kind falls back to the estimate + return type(e).__name__ + if hashlib.sha256(data).hexdigest() != source.sha256: + return "SHA-256 mismatch" + rows: Final = (line.split() for line in data.splitlines() if line) + return tiktoken.Encoding( + name="kimi_k3", + pat_str=_PATTERN, + mergeable_ranks={base64.b64decode(token): int(rank) for token, rank in rows}, + special_tokens={}, + ) + + +def load_kimi_k3_encoding(source: KimiK3TokenizerSource) -> tiktoken.Encoding | None: + loaded: Final = _loaded(source) + if isinstance(loaded, tiktoken.Encoding): + return loaded + verbose_logger.warning( + "Kimi K3 tokenizer unavailable (%s), so Kimi K3 token counts are estimated until a later count loads it. " + "Allow downloads of %s from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache", + loaded, + KIMI_K3_RANKS_FILE, + KIMI_K3_REPO, + KIMI_K3_REVISION, + ) + return None + + +def _started_load(source: KimiK3TokenizerSource) -> "Future[tiktoken.Encoding | None]": + encoding: Final = Future[tiktoken.Encoding | None]() + threading.Thread( + target=lambda: encoding.set_result(load_kimi_k3_encoding(source)), name="litellm-kimi-k3-tokenizer", daemon=True + ).start() + return encoding + + +@dataclass(frozen=True, slots=True) +class _EncodingLoad: + encoding: "Future[tiktoken.Encoding | None]" + started_at: float + is_retry: bool + + def retry_due(self, now: float, retry_after_seconds: float) -> bool: + return self.encoding.done() and self.encoding.result() is None and now - self.started_at >= retry_after_seconds + + +@dataclass(slots=True) +class KimiK3Encodings: + source: KimiK3TokenizerSource + retry_after_seconds: float + clock: Callable[[], float] = time.monotonic + current: _EncodingLoad | None = None + + def _load(self) -> _EncodingLoad: + now: Final = self.clock() + current: Final = self.current + if current is not None and not current.retry_due(now, self.retry_after_seconds): + return current + load: Final = _EncodingLoad(encoding=_started_load(self.source), started_at=now, is_retry=current is not None) + self.current = load + return load + + async def encoding(self, wait_seconds: float) -> tiktoken.Encoding | None: + load: Final = self._load() + remaining: Final = 0 if load.is_retry else load.started_at + wait_seconds - self.clock() + if load.encoding.done() or remaining <= 0: + return load.encoding.result() if load.encoding.done() else None + await asyncio.wait((asyncio.wrap_future(load.encoding),), timeout=remaining) + if load.encoding.done(): + return load.encoding.result() + verbose_logger.warning( + "Kimi K3 tokenizer still loading after %s s, so Kimi K3 token counts are estimated until it loads. " + "Check that %s can be downloaded from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache", + wait_seconds, + KIMI_K3_RANKS_FILE, + KIMI_K3_REPO, + KIMI_K3_REVISION, + ) + return None + + +class _Converse(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class _Text(_Converse): + text: str + + +class _CachePoint(_Converse): + cache_point: JsonValue = Field(alias="cachePoint") + + +class _ToolUse(_Converse): + tool_use_id: str = Field(alias="toolUseId") + name: str + arguments: Mapping[str, JsonValue] = Field(alias="input") + + +class _ToolUseBlock(_Converse): + tool_use: _ToolUse = Field(alias="toolUse") + + +class _ToolResult(_Converse): + tool_use_id: str = Field(alias="toolUseId") + content: tuple[_Text, ...] + status: str | None = None + + +class _ToolResultBlock(_Converse): + tool_result: _ToolResult = Field(alias="toolResult") + + +class _ReasoningText(_Converse): + text: str + signature: str | None = None + + +class _Reasoning(_Converse): + reasoning_text: _ReasoningText = Field(alias="reasoningText") + + +class _ReasoningBlock(_Converse): + reasoning_content: _Reasoning = Field(alias="reasoningContent") + + +class _UserMessage(_Converse): + role: Literal["user"] + content: tuple[_Text | _ToolResultBlock | _CachePoint, ...] + + +class _AssistantMessage(_Converse): + role: Literal["assistant"] + content: tuple[_Text | _ToolUseBlock | _ReasoningBlock | _CachePoint, ...] + + def calls(self) -> tuple[_ToolUse, ...]: + return tuple(block.tool_use for block in self.content if isinstance(block, _ToolUseBlock)) + + +class _InputSchema(_Converse): + definition: JsonValue = Field(alias="json") + + +class _ToolSpec(_Converse): + name: str + description: str | None = None + input_schema: _InputSchema = Field(alias="inputSchema") + + def declared(self) -> Mapping[str, JsonValue]: + described: Final[Mapping[str, JsonValue]] = ( + {} if self.description is None else {"description": self.description} + ) + return {"name": self.name, **described, "parameters": self.input_schema.definition} + + +class _ToolSpecBlock(_Converse): + tool_spec: _ToolSpec = Field(alias="toolSpec") + + +class _ToolConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + tools: tuple[_ToolSpecBlock | _CachePoint, ...] + + +_Message: TypeAlias = Annotated[_UserMessage | _AssistantMessage, Field(discriminator="role")] + + +def _result_ids(message: _Message) -> frozenset[str]: + if isinstance(message, _AssistantMessage): + return frozenset() + return frozenset(block.tool_result.tool_use_id for block in message.content if isinstance(block, _ToolResultBlock)) + + +def _call_ids(message: _Message | None) -> frozenset[str]: + if not isinstance(message, _AssistantMessage): + return frozenset() + return frozenset(call.tool_use_id for call in message.calls()) + + +class _ConverseRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + messages: tuple[_Message, ...] + system: tuple[_Text | _CachePoint, ...] = () + tool_config: _ToolConfig | None = Field(default=None, alias="toolConfig") + + @model_validator(mode="after") + def _results_answer_the_previous_calls(self) -> Self: + pairs: Final = zip((None, *self.messages), self.messages, strict=False) + if not all(_result_ids(message) <= _call_ids(previous) for previous, message in pairs): + raise ValueError("a toolResult does not answer a toolUse of the previous assistant message") + return self + + +class _Control(Enum): + OPEN = "<|open|>" + CLOSE = "<|close|>" + SEP = "<|sep|>" + END_OF_MSG = "<|end_of_msg|>" + + +_Segment: TypeAlias = str | _Control +_Segments: TypeAlias = tuple[_Segment, ...] + + +def _open(tag: str, attrs: Iterable[tuple[str, str]] = ()) -> _Segments: + escaped: Final = ( + (f" {key}", '="', value.replace("&", "&").replace('"', """), '"') for key, value in attrs + ) + return (_Control.OPEN, tag, *chain.from_iterable(escaped), _Control.SEP) + + +def _close(tag: str) -> _Segments: + return (_Control.CLOSE, tag, _Control.SEP) + + +def _message(attrs: Iterable[tuple[str, str]], body: Iterable[_Segment]) -> _Segments: + return (*_open("message", attrs), *body, *_close("message"), _Control.END_OF_MSG) + + +def _tool_declare(specs: Sequence[_ToolSpec]) -> _Segments: + tools: Final = json.dumps( + [{"type": "function", "function": spec.declared()} for spec in specs], + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + return _message( + (("role", "system"), ("type", "tool-declare")), + (f"# Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n{tools}\n```",), + ) + + +def _text(blocks: Iterable[_Converse]) -> str: + return "".join(block.text for block in blocks if isinstance(block, _Text)) + + +def _json_type(value: JsonValue) -> str: + match value: + case bool(): + return "boolean" + case None: + return "null" + case int() | float(): + return "number" + case str(): + return "string" + case dict(): + return "object" + case _: + return "array" + + +def _argument(key: str, value: JsonValue) -> _Segments: + text: Final = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False) + return (*_open("argument", (("key", key), ("type", _json_type(value)))), text, *_close("argument")) + + +def _call(index: int, call: _ToolUse) -> _Segments: + arguments: Final = chain.from_iterable(_argument(key, value) for key, value in call.arguments.items()) + return (*_open("call", (("tool", call.name), ("index", str(index)))), *arguments, *_close("call")) + + +def _tool_calls(calls: Sequence[_ToolUse]) -> _Segments: + if not calls: + return () + rendered: Final = chain.from_iterable(_call(index, call) for index, call in enumerate(calls, start=1)) + return (*_open("tools"), *rendered, *_close("tools")) + + +def _assistant_message(message: _AssistantMessage) -> _Segments: + reasoning: Final = "".join( + block.reasoning_content.reasoning_text.text for block in message.content if isinstance(block, _ReasoningBlock) + ) + return _message( + (("role", "assistant"),), + ( + *_open("think"), + reasoning, + *_close("think"), + *_open("response"), + _text(message.content), + *_close("response"), + *_tool_calls(message.calls()), + ), + ) + + +def _tool_messages(results: Sequence[_ToolResult], previous: _AssistantMessage) -> _Segments: + names: Final = {call.tool_use_id: call.name for call in previous.calls()} + return tuple( + chain.from_iterable( + _message( + (("role", "tool"), ("tool", names[result.tool_use_id]), ("index", str(index))), + (_text(result.content),), + ) + for index, result in enumerate(results, start=1) + ) + ) + + +def _user_run(blocks: Sequence[_Text | _ToolResultBlock], previous: _Message | None) -> _Segments: + results: Final = tuple(block.tool_result for block in blocks if isinstance(block, _ToolResultBlock)) + if results and isinstance(previous, _AssistantMessage): + return _tool_messages(results, previous) + return _message((("role", "user"),), (_text(blocks),)) + + +def _user_messages(message: _UserMessage, previous: _Message | None) -> _Segments: + blocks: Final = tuple(block for block in message.content if not isinstance(block, _CachePoint)) + runs: Final = groupby(blocks, key=lambda block: isinstance(block, _ToolResultBlock)) + return tuple(chain.from_iterable(_user_run(tuple(run), previous) for _, run in runs)) + + +def _turn(previous: _Message | None, message: _Message) -> _Segments: + if isinstance(message, _AssistantMessage): + return _assistant_message(message) + return _user_messages(message, previous) + + +def _rendered(request: _ConverseRequest) -> _Segments: + specs: Final = ( + () + if request.tool_config is None + else tuple(block.tool_spec for block in request.tool_config.tools if isinstance(block, _ToolSpecBlock)) + ) + system: Final = _text(request.system) + turns: Final = chain.from_iterable( + _turn(previous, message) for previous, message in zip((None, *request.messages), request.messages, strict=False) + ) + return ( + *(_tool_declare(specs) if specs else ()), + *(_message((("role", "system"),), (system,)) if system else ()), + *turns, + *_open("message", (("role", "assistant"),)), + *_open("think"), + ) + + +def _text_tokens(encoding: tiktoken.Encoding, text: str) -> int: + spans: Final = (match.span() for match in _CHUNK.finditer(text)) + return sum(len(encoding.encode_ordinary(text[start:end])) for start, end in spans) + + +class _TextPart(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + type: Literal["text"] + + +class _TextOnlyMessage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + content: str | tuple[_TextPart, ...] | None = None + + +def _text_only(message: AllMessageValues) -> bool: + try: + _TextOnlyMessage.model_validate(message) + except ValidationError: + return False + return True + + +def _as_converse_input(tool: ChatCompletionToolParam) -> tuple[Mapping[str, object], ...]: + return chat_completion_tools_for(tool) if is_flat_responses_tool(tool) else (tool,) + + +def _converse_request( + model: str, messages: Sequence[AllMessageValues], tools: Sequence[ChatCompletionToolParam] | None +) -> _ConverseRequest | None: + if not all(_text_only(message) for message in messages): + return None + try: + chat_tools: Final = tuple( + chat_tool for tool in tools or () for chat_tool in _as_converse_input(tool) + ) # comprehension-ok: flattens each tool's translation + return _ConverseRequest.model_validate( + litellm.AmazonConverseConfig().transform_request( + model=model, + messages=copy.deepcopy(list(messages)), + optional_params={"tools": copy.deepcopy(list(chat_tools))} if chat_tools else {}, + litellm_params={}, + headers={}, + ) + ) + except Exception: # noqa: BLE001 # a body litellm cannot send to Bedrock, or cannot render, falls back to the estimate + return None + + +def kimi_k3_tokens( + encoding: tiktoken.Encoding, + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam] | None, +) -> int | None: + request: Final = _converse_request(model, messages, tools) + if request is None: + return None + return sum( + 1 if isinstance(segment, _Control) else _text_tokens(encoding, segment) for segment in _rendered(request) + ) + + +def kimi_k3_anthropic_tokens(encoding: tiktoken.Encoding, model: str, fields: Mapping[str, object]) -> int | None: + body: Final = {**fields, "model": model, "max_tokens": 1} + try: + translated, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + copy.deepcopy(body), # pyright: ignore[reportArgumentType] # raw /v1/messages fields + custom_llm_provider="bedrock", + ) + except Exception: # noqa: BLE001 # a request the adapter cannot translate falls back to the estimate + return None + match translated: + case {"tools": tools}: + return kimi_k3_tokens(encoding, model, translated["messages"], tools) + case _: + return kimi_k3_tokens(encoding, model, translated["messages"], None) + + +def kimi_k3_template_tokens(message_count: int, has_tools: bool) -> int: + return 14 + 13 * max(message_count - 1, 0) + (40 if has_tools else 0) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 92b3b85dc51..3d04954dd41 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence from typing import Any, Final import httpx -from pydantic import TypeAdapter +from pydantic import JsonValue, TypeAdapter import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -14,7 +14,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = { "1:1": 1 / 1, @@ -495,21 +495,54 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: import copy from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter + from litellm.llms.vertex_ai.count_tokens.context import ( + VertexGeminiCountContextError, + count_context_error_response, + gemini_count_contents, + inference_request_view, + provider_error_response, + unconvertible_messages_error, + vertex_gemini_count_context, + ) + + request_view: Final = inference_request_view(messages or [], system, messages_format) + count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) + if isinstance(count_context, VertexGeminiCountContextError): + return count_context_error_response(count_context, request_model, model_to_use) + blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format) + if blocks_error is not None: + return count_context_error_response(blocks_error, request_model, model_to_use) deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) + resolved_contents: Final = ( + contents + if contents is not None + else await gemini_count_contents( + request_view.messages, messages_format, model_to_use, "gemini", count_tokens_params_request + ) + ) + if isinstance(resolved_contents, VertexGeminiCountContextError): + return count_context_error_response(resolved_contents, request_model, model_to_use) count_tokens_params: Final = { "model": model_to_use, - "contents": contents, + "contents": resolved_contents, + "system_instruction": count_context.system_instruction, + "tools": count_context.tools, } count_tokens_params_request.update(count_tokens_params) - result: Final = await GoogleAIStudioTokenCounter().acount_tokens( - **count_tokens_params_request, - ) + try: + result: Final = await GoogleAIStudioTokenCounter().acount_tokens( + **count_tokens_params_request, + ) + except litellm.APIError as error: + return provider_error_response(error, request_model, model_to_use) if result is not None: return TokenCountResponse( diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py index 67e84f6da24..46615534747 100644 --- a/litellm/llms/gemini/count_tokens/handler.py +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -1,9 +1,10 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx import litellm -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.types.utils import LlmProviders if TYPE_CHECKING: @@ -28,12 +29,12 @@ class GoogleAIStudioTokenCounter: """ import copy - from google.genai.types import FunctionResponse - # Handle None or empty contents if not contents: return contents + from google.genai.types import FunctionResponse + cleaned_contents: Final = copy.deepcopy(contents) for content in cleaned_contents: @@ -54,6 +55,28 @@ class GoogleAIStudioTokenCounter: base_url: Final = api_base or "https://generativelanguage.googleapis.com" return f"{base_url}/v1beta/models/{model}:countTokens" + def __init__(self, client: AsyncHTTPHandler | None = None) -> None: + super().__init__() + self._client: Final = client + + def _build_count_tokens_request_body( + self, + model: str, + contents: object, + system_instruction: Mapping[str, object] | None, + tools: Sequence[Mapping[str, object]] | None, + ) -> Mapping[str, object]: + if system_instruction is None and not tools: + return {"contents": contents} + + generate_content_request: Final = { + "model": model if model.startswith("models/") else f"models/{model}", + "contents": contents, + **({"systemInstruction": dict(system_instruction)} if system_instruction is not None else {}), + **({"tools": [dict(tool) for tool in tools]} if tools else {}), + } + return {"generateContentRequest": generate_content_request} + async def validate_environment( self, api_base: str | None = None, @@ -84,6 +107,8 @@ class GoogleAIStudioTokenCounter: api_key: str | None = None, api_base: str | None = None, timeout: float | httpx.Timeout | None = None, + system_instruction: Mapping[str, object] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, **kwargs: object, ) -> dict[str, Any]: """ @@ -130,14 +155,19 @@ class GoogleAIStudioTokenCounter: # Prepare request body - clean up contents to remove unsupported fields cleaned_contents: Final = self._clean_contents_for_gemini_api(contents) - request_body: Final = {"contents": cleaned_contents} + request_body: Final = self._build_count_tokens_request_body( + model=model, + contents=cleaned_contents, + system_instruction=system_instruction, + tools=tools, + ) - async_httpx_client: Final = get_async_httpx_client( - llm_provider=LlmProviders.GEMINI, + async_httpx_client: Final = ( + self._client if self._client is not None else get_async_httpx_client(llm_provider=LlmProviders.GEMINI) ) try: - response: Final = await async_httpx_client.post(url=url, headers=headers, json=request_body) + response: Final = await async_httpx_client.post(url=url, headers=headers, json=dict(request_body)) # Check for HTTP errors response.raise_for_status() diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 0d7a553d633..6b34f04a36a 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -5,7 +5,7 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint. """ import json -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Final import httpx @@ -31,9 +31,10 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): api_key: str, api_base: str | None = None, timeout: float | httpx.Timeout | None = None, - tools: list[dict[str, object]] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, - ) -> dict[str, Any]: + text: Mapping[str, object] | None = None, + ) -> dict[str, object]: """ Handle a token counting request to OpenAI's Responses API. @@ -53,6 +54,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): input=input, tools=tools, instructions=instructions, + text=text, ) endpoint_url: Final = self.get_openai_count_tokens_endpoint(api_base) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index c05d943b7cf..cd0c3c4ec86 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -2,24 +2,47 @@ OpenAI Token Counter implementation using the Responses API /input_tokens endpoint. """ -import os +from collections.abc import Mapping, Sequence from typing import Any, Final +import openai +from pydantic import JsonValue +from typing_extensions import assert_never + from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify +from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.common_utils import OpenAIError, get_openai_credentials from litellm.llms.openai.responses.count_tokens.handler import ( OpenAICountTokensHandler, ) from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, + ResponsesCountRequestParts, ) -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.llms.openai.workload_identity import ( + OpenAIWorkloadIdentityConfig, + get_workload_identity_bearer_token, + resolve_openai_workload_identity_config, +) +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests openai_count_tokens_handler: Final = OpenAICountTokensHandler() +def _workload_identity_config(litellm_params: Mapping[str, object]) -> OpenAIWorkloadIdentityConfig | None: + api_base: Final = litellm_params.get("api_base") + return resolve_openai_workload_identity_config( + api_key=None, + api_base=api_base if isinstance(api_base, str) else None, + ) + + class OpenAITokenCounter(BaseTokenCounter): """Token counter implementation for OpenAI provider using the Responses API.""" @@ -29,6 +52,37 @@ class OpenAITokenCounter(BaseTokenCounter): ) -> bool: return custom_llm_provider == LlmProviders.OPENAI.value + @staticmethod + def _system_instructions(system: object) -> str | None: + if system is None or isinstance(system, str): + return system + return LiteLLMAnthropicToResponsesAPIAdapter.system_blocks_to_instructions(system) or None + + @staticmethod + def _responses_request_parts( + model_to_use: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + system: object, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue, + ) -> ResponsesCountRequestParts: + match messages_format: + case "anthropic": + return OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model=model_to_use, messages=messages, system=system, tools=tools, output_format=output_format + ) + case "openai": + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( + [dict(message) for message in messages] + ) + fallback_instructions: Final = OpenAITokenCounter._system_instructions(system) + return OpenAICountTokensConfig.count_request_parts( + input_items, instructions if instructions is not None else fallback_instructions, tools + ) + case _: + assert_never(messages_format) + async def count_tokens( self, model_to_use: str, @@ -38,6 +92,8 @@ class OpenAITokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using OpenAI's Responses API /input_tokens endpoint. @@ -45,39 +101,76 @@ class OpenAITokenCounter(BaseTokenCounter): if not messages: return None + malformed_error: Final = ( + OpenAICountTokensConfig.malformed_anthropic_messages_error(messages) + if messages_format == "anthropic" + else None + ) + request_error: Final = malformed_error or unconvertible_blocks_error(messages, messages_format) + if request_error is not None: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=request_error, + status_code=400, + rejected_before_provider_call=True, + ) + deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) - # Get OpenAI API key from deployment config or environment - api_key = litellm_params.get("api_key") - if not api_key: - api_key = os.getenv("OPENAI_API_KEY") + credentials: Final = get_openai_credentials( + api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key") + ) + credential: Final = credentials.api_key or _workload_identity_config(litellm_params) - if not api_key: - verbose_logger.warning("No OpenAI API key found for token counting") - return None + if credential is None: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=( + f"No OpenAI API key is set for model {request_model or model_to_use}, so tokens can't be " + "counted. Set api_key or workload identity on the deployment, or OPENAI_API_KEY in the proxy " + "environment." + ), + status_code=401, + ) - api_base: Final = litellm_params.get("api_base") + api_base: Final = credentials.api_base - # Convert chat messages to Responses API input format - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) - - # Use system param if instructions not extracted from messages - if instructions is None and system is not None: - instructions = system if isinstance(system, str) else str(system) + input_items, instructions, counted_tools, text = self._responses_request_parts( + model_to_use=model_to_use, + messages=messages, + tools=tools, + system=system, + messages_format=messages_format, + output_format=output_format, + ) # If no input items were produced (e.g., system-only messages), fall back to local counting if not input_items: return None try: + api_key: Final = ( + credential + if isinstance(credential, str) + else await asyncify(get_workload_identity_bearer_token)(credential) + ) result: Final = await openai_count_tokens_handler.handle_count_tokens_request( model=model_to_use, - input=input_items if input_items is not None else [], + input=input_items, api_key=api_key, api_base=api_base, - tools=tools, + tools=counted_tools, instructions=instructions, + text=text, ) if result is not None: @@ -99,6 +192,17 @@ class OpenAITokenCounter(BaseTokenCounter): error_message=e.message, status_code=e.status_code, ) + except openai.AuthenticationError as e: + verbose_logger.warning("OpenAI workload identity token exchange failed: %s", e.message) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=e.message, + status_code=401, + ) except Exception as e: verbose_logger.warning("Error calling OpenAI CountTokens API: %s", e) return TokenCountResponse( diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 9f596505f91..eda9777801c 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -5,10 +5,18 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp """ from collections.abc import Mapping, Sequence -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypeAlias +from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict +from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema + class ResponsesInputTextPart(TypedDict): type: ReadOnly[Literal["input_text"]] @@ -31,6 +39,30 @@ ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart | Response ResponsesContentRole = Literal["user", "assistant"] +ResponsesCountRequestParts: TypeAlias = tuple[ + tuple[Mapping[str, object], ...], + str | None, + tuple[Mapping[str, object], ...] | None, + Mapping[str, object] | None, +] +_RESPONSES_COUNT_REQUEST_PARTS: Final = TypeAdapter(ResponsesCountRequestParts) +_TRANSLATED_REQUEST: Final = TypeAdapter(Mapping[str, object]) + + +class _AnthropicCountMessage(TypedDict): + role: ReadOnly[str] + + +_ANTHROPIC_COUNT_MESSAGES: Final = TypeAdapter(tuple[_AnthropicCountMessage, ...]) +_ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema) + + +def _anthropic_output_schema(output_format: JsonValue) -> AnthropicOutputSchema | None: + try: + return _ANTHROPIC_OUTPUT_SCHEMA.validate_python(output_format) + except ValidationError: + return None + 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 @@ -118,8 +150,9 @@ class OpenAICountTokensConfig: self, model: str, input: str | Sequence[object], - tools: list[dict[str, Any]] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, + text: Mapping[str, object] | None = None, ) -> dict[str, object]: """ Transform request to OpenAI Responses API token counting format. @@ -135,7 +168,10 @@ class OpenAICountTokensConfig: request["instructions"] = instructions if tools is not None: - request["tools"] = self._transform_tools_for_responses_api(tools) + request["tools"] = self._transform_tools_for_responses_api([dict(tool) for tool in tools]) + + if text is not None: + request["text"] = text return request @@ -183,7 +219,7 @@ class OpenAICountTokensConfig: @staticmethod def messages_to_responses_input( messages: list[dict[str, Any]], - ) -> tuple: + ) -> tuple[list[dict[str, object]], str | None]: """ Convert standard chat messages format to OpenAI Responses API input format. @@ -246,3 +282,47 @@ class OpenAICountTokensConfig: instructions: Final = "\n".join(instructions_parts) if instructions_parts else None return input_items, instructions + + @staticmethod + def malformed_anthropic_messages_error(messages: Sequence[Mapping[str, object]]) -> str | None: + try: + _ANTHROPIC_COUNT_MESSAGES.validate_python(messages) + except ValidationError as e: + return "; ".join( + f"messages.{'.'.join(str(part) for part in error['loc'])}: {error['msg']}" for error in e.errors() + ) + return None + + @staticmethod + def anthropic_messages_to_responses_input( + model: str, + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + output_format: JsonValue = None, + ) -> ResponsesCountRequestParts: + """Convert with the adapter /v1/messages inference uses, so the count sees the same input it bills.""" + anthropic_request: Final = AnthropicMessagesRequest( + model=model, + messages=normalize_native_messages(messages), + system=system if isinstance(system, (str, list)) else None, + tools=[dict(tool) for tool in tools] if tools else None, + output_format=_anthropic_output_schema(output_format), + ) + translated: Final = _TRANSLATED_REQUEST.validate_python( + LiteLLMAnthropicToResponsesAPIAdapter().translate_request( + anthropic_request, include_encrypted_reasoning=False + ) + ) + return OpenAICountTokensConfig.count_request_parts( + ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(translated.get("input")), # pyright: ignore[reportPrivateUsage] # the input restore inference runs + translated.get("instructions"), + translated.get("tools"), + translated.get("text"), + ) + + @staticmethod + def count_request_parts( + input_items: object, instructions: object, tools: object, text: object = None + ) -> ResponsesCountRequestParts: + return _RESPONSES_COUNT_REQUEST_PARTS.validate_python((input_items, instructions, tools, text)) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 6b759065706..f1afe04719a 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -7,7 +7,7 @@ from types import MappingProxyType from typing import Any, Final, Literal, cast, get_type_hints import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -22,7 +22,7 @@ from litellm.types.llms.vertex_ai import ( PartType, Schema, ) -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType( @@ -1271,6 +1271,8 @@ class VertexAITokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: import copy @@ -1342,30 +1344,55 @@ class VertexAITokenCounter(BaseTokenCounter): original_response=result, ) else: - from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter - from litellm.llms.vertex_ai.gemini.transformation import ( - gemini_convert_messages_with_history, # pyright: ignore[reportPrivateUsage] # shared helper already used by gemini/chat, context_caching, and vertex_and_google_ai_studio_gemini + from litellm.llms.vertex_ai.count_tokens.context import ( + VertexGeminiCountContextError, + count_context_error_response, + gemini_count_contents, + inference_request_view, + provider_error_response, + unconvertible_messages_error, + vertex_gemini_count_context, ) + from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter + + request_view: Final = inference_request_view(messages or [], system, messages_format) + count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) + if isinstance(count_context, VertexGeminiCountContextError): + return count_context_error_response(count_context, request_model, model_to_use) + blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format) + if blocks_error is not None: + return count_context_error_response(blocks_error, request_model, model_to_use) resolved_contents: Final = ( - contents if contents is not None else gemini_convert_messages_with_history(messages=messages or []) + contents + if contents is not None + else await gemini_count_contents( + request_view.messages, messages_format, model_to_use, "vertex_ai", count_tokens_params_request + ) ) + if isinstance(resolved_contents, VertexGeminiCountContextError): + return count_context_error_response(resolved_contents, request_model, model_to_use) count_tokens_params: Final = { "model": model_to_use, "contents": resolved_contents, + "system_instruction": count_context.system_instruction, + "tools": count_context.tools, } count_tokens_params_request.update(count_tokens_params) - result = await VertexAITokenCounter().acount_tokens( - **count_tokens_params_request, - ) + try: + result = await VertexAITokenCounter().acount_tokens( + **count_tokens_params_request, + ) + except litellm.APIError as error: + return provider_error_response(error, request_model, model_to_use) - if result is not None and "totalTokens" in result: + if "totalTokens" in result: return TokenCountResponse( total_tokens=result["totalTokens"], request_model=request_model, model_used=model_to_use, - tokenizer_type=result.get("tokenizer_used", ""), + tokenizer_type=result.get("tokenizer_used", "gemini"), original_response=result, ) diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py new file mode 100644 index 00000000000..27f2987ea50 --- /dev/null +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -0,0 +1,242 @@ +import copy +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal + +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for +from litellm.types.llms.anthropic import AnthropicMessagesTextParam +from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.vertex_ai import ContentType +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse + +if TYPE_CHECKING: + from litellm.exceptions import APIError + from litellm.llms.anthropic.pass_through.context_management.result import PolyfillResult + + +@dataclass(frozen=True, slots=True) +class VertexGeminiCountContext: + system_instruction: Mapping[str, object] | None + tools: tuple[Mapping[str, object], ...] | None + + +@dataclass(frozen=True, slots=True) +class VertexGeminiCountContextError: + error_message: str + + +_SYSTEM_BLOCKS: Final = TypeAdapter(tuple[AnthropicMessagesTextParam, ...]) +_TOOLS: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) +_VERTEX_TOOL_KEYS: Final = MappingProxyType( + { + "function_declarations": "functionDeclarations", + "url_context": "urlContext", + "code_execution": "codeExecution", + } +) + + +def _system_parts( + system: object | None, +) -> tuple[Mapping[str, JsonValue], ...] | VertexGeminiCountContextError: + if system is None: + return () + if isinstance(system, str): + return ({"text": system},) if system else () + try: + blocks: Final = _SYSTEM_BLOCKS.validate_python(system) + except ValidationError: + return VertexGeminiCountContextError(error_message="Gemini system instructions only support text blocks.") + return tuple({"text": block["text"]} for block in blocks) + + +_INVALID_TOOLS: Final = VertexGeminiCountContextError( + error_message="tools must be a list of objects, and each function tool needs a JSON object schema." +) + + +def _has_object_schema(tool: Mapping[str, JsonValue]) -> bool: + function: Final = tool.get("function") + schema: Final = ( + function.get("parameters") if isinstance(function, dict) else tool.get("input_schema", tool.get("parameters")) + ) + return schema is None or isinstance(schema, dict) + + +def _is_function_tool(tool: Mapping[str, JsonValue]) -> bool: + return isinstance(tool.get("name"), str) or "input_schema" in tool or "function" in tool + + +def _map_function_tools( + tools: object | None, +) -> tuple[Mapping[str, object], ...] | None | VertexGeminiCountContextError: + try: + tool_values: Final = _TOOLS.validate_python(tools or ()) + except ValidationError: + return _INVALID_TOOLS + function_tools: Final = tuple(tool for tool in tool_values if _is_function_tool(tool)) + if not function_tools: + return None + if not all(_has_object_schema(tool) for tool in function_tools): + return _INVALID_TOOLS + + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + chat_tools: Final = tuple( + chat_tool for tool in function_tools for chat_tool in chat_completion_tools_for(tool) + ) # comprehension-ok: flattens each tool's translation + try: + mapped_tools: Final = VertexGeminiConfig()._map_function( # pyright: ignore[reportPrivateUsage] # the tool mapping inference uses + value=[dict(tool) for tool in chat_tools], + optional_params={}, + ) + except (AttributeError, TypeError, ValueError): + return _INVALID_TOOLS + return tuple( + MappingProxyType({_VERTEX_TOOL_KEYS.get(key, key): value for key, value in tool.items()}) + for tool in mapped_tools + ) + + +def vertex_gemini_count_context( + system: object | None, + tools: object | None, +) -> VertexGeminiCountContext | VertexGeminiCountContextError: + system_parts: Final = _system_parts(system) + if isinstance(system_parts, VertexGeminiCountContextError): + return system_parts + function_tools: Final = _map_function_tools(tools) + if isinstance(function_tools, VertexGeminiCountContextError): + return function_tools + return VertexGeminiCountContext( + system_instruction=MappingProxyType({"parts": system_parts}) if system_parts else None, + tools=function_tools, + ) + + +def unconvertible_messages_error( + messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat +) -> VertexGeminiCountContextError | None: + from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error + + error_message: Final = unconvertible_blocks_error(messages, messages_format) + return None if error_message is None else VertexGeminiCountContextError(error_message=error_message) + + +def count_context_error_response( + error: VertexGeminiCountContextError, request_model: str, model_used: str +) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_used, + tokenizer_type="gemini", + error=True, + error_message=error.error_message, + status_code=400, + rejected_before_provider_call=True, + ) + + +def provider_error_response(error: "APIError", request_model: str, model_used: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_used, + tokenizer_type="gemini", + error=True, + error_message=error.message, + status_code=error.status_code, + ) + + +@dataclass(frozen=True, slots=True) +class InferenceRequestView: + messages: Sequence[Mapping[str, object]] + system: object | None + + +def _client_compaction(messages: Sequence[Mapping[str, object]], system: object | None) -> "PolyfillResult | None": + from litellm.llms.anthropic.pass_through.context_management.editors.compact import ( + apply_client_compaction_block_history, + ) + + chat_messages: Final = [dict(message) for message in messages] + if system is None or isinstance(system, str): + return apply_client_compaction_block_history(messages=chat_messages, system=system) + try: + blocks: Final = _SYSTEM_BLOCKS.validate_python(system) + except ValidationError: + return None + return apply_client_compaction_block_history(messages=chat_messages, system=[dict(block) for block in blocks]) + + +def inference_request_view( + messages: Sequence[Mapping[str, object]], system: object | None, messages_format: TokenCountMessagesFormat +) -> InferenceRequestView: + """Apply the history rewrites /v1/messages runs before it converts a body for Gemini.""" + from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages + + match messages_format: + case "openai": + return InferenceRequestView(messages=messages, system=system) + case "anthropic": + normalized: Final = normalize_native_messages(messages) + compacted: Final = _client_compaction(normalized, system) + if compacted is None: + return InferenceRequestView(messages=normalized, system=system) + return InferenceRequestView(messages=compacted.messages, system=compacted.system) + case _: + assert_never(messages_format) + + +def _chat_messages( + messages: Sequence[Mapping[str, object]], + messages_format: TokenCountMessagesFormat, + model: str, + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> Sequence[AllMessageValues]: + from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter + + match messages_format: + case "anthropic": + return LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=list(messages), # pyright: ignore[reportArgumentType] # raw Anthropic body, read the way /v1/messages inference reads it + model=model, + custom_llm_provider=custom_llm_provider, + ) + case "openai": + return copy.deepcopy(list(messages)) # pyright: ignore[reportReturnType] # chat messages from the request, copied because the AI Studio conversion rewrites them in place + case _: + assert_never(messages_format) + + +async def gemini_count_contents( + messages: Sequence[Mapping[str, object]], + messages_format: TokenCountMessagesFormat, + model: str, + custom_llm_provider: Literal["vertex_ai", "gemini"], + litellm_params: Mapping[str, object], +) -> Sequence[ContentType] | VertexGeminiCountContextError: + from litellm.exceptions import BadRequestError + from litellm.litellm_core_utils.asyncify import asyncify + from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media + from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig + from litellm.llms.vertex_ai.gemini.transformation import should_inline_remote_media + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + config: Final = GoogleAIStudioGeminiConfig() if custom_llm_provider == "gemini" else VertexGeminiConfig() + try: + chat_messages: Final = await async_inline_remote_media( + list(_chat_messages(messages, messages_format, model, custom_llm_provider)), + should_inline=should_inline_remote_media(custom_llm_provider), + ) + return await asyncify(config._transform_messages)( # pyright: ignore[reportPrivateUsage] # the message transform inference runs + messages=chat_messages, model=model, litellm_params=dict(litellm_params) + ) + except BadRequestError as error: + return VertexGeminiCountContextError(error_message=error.message) diff --git a/litellm/llms/vertex_ai/count_tokens/handler.py b/litellm/llms/vertex_ai/count_tokens/handler.py index 47a08ff054d..03970099f2e 100644 --- a/litellm/llms/vertex_ai/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/count_tokens/handler.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping, Sequence from typing import Any, Final from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter @@ -5,6 +6,19 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase): + def _build_count_tokens_request_body( + self, + model: str, + contents: object, + system_instruction: Mapping[str, object] | None, + tools: Sequence[Mapping[str, object]] | None, + ) -> Mapping[str, object]: + return { + "contents": contents, + **({"systemInstruction": dict(system_instruction)} if system_instruction else {}), + **({"tools": [dict(tool) for tool in tools]} if tools else {}), + } + async def validate_environment( self, api_base: str | None = None, @@ -16,29 +30,29 @@ class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase): """ Returns a Tuple of headers and url for the Vertex AI countTokens endpoint. """ - litellm_params = litellm_params or {} - vertex_credentials: Final = self.get_vertex_ai_credentials(litellm_params=litellm_params) - vertex_project = self.get_vertex_ai_project(litellm_params=litellm_params) - vertex_location: Final = self.get_vertex_ai_location(litellm_params=litellm_params) - _auth_header, vertex_project = await self._ensure_access_token_async( + resolved_litellm_params: Final = litellm_params or {} + vertex_credentials: Final = self.get_vertex_ai_credentials(litellm_params=resolved_litellm_params) + vertex_project: Final = self.get_vertex_ai_project(litellm_params=resolved_litellm_params) + vertex_location: Final = self.get_vertex_ai_location(litellm_params=resolved_litellm_params) + _auth_header, resolved_vertex_project = await self._ensure_access_token_async( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", ) - auth_header, api_base = self._get_token_and_url( + auth_header, url = self._get_token_and_url( model=model, gemini_api_key=None, auth_header=_auth_header, - vertex_project=vertex_project, + vertex_project=resolved_vertex_project, vertex_location=vertex_location, vertex_credentials=vertex_credentials, stream=False, custom_llm_provider="vertex_ai", - api_base=None, + api_base=api_base, mode="count_tokens", ) headers = { "Authorization": f"Bearer {auth_header}", } - return headers, api_base + return headers, url diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 15e5e440b91..7af9a0790c0 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -6,7 +6,7 @@ Why separate file? Make it easy to see how transformation works import os import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote @@ -1311,6 +1311,10 @@ def _vertex_inlines(media: RemoteMedia) -> bool: ) +def should_inline_remote_media(custom_llm_provider: str) -> Callable[[RemoteMedia], bool]: + return _ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + + async def async_transform_request_body( gemini_api_key: str | None, messages: list[AllMessageValues], @@ -1353,7 +1357,7 @@ async def async_transform_request_body( ) inlined_messages: Final = await async_inline_remote_media( - messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + messages, should_inline=should_inline_remote_media(custom_llm_provider) ) if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8d0157f0b97..9de0a3d967f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -20077,7 +20077,13 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "token_count_framing": 97 + }, + "databricks/system.ai.deepseek-v4-flash-0731": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 97 }, "databricks/databricks-deepseek-v4-pro-0813": { "cache_creation_input_token_cost": 1.31999e-06, @@ -21068,7 +21074,13 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 100 + }, + "databricks/system.ai.kimi-k3": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 100 }, "databricks/databricks-llama-4-maverick": { "cache_creation_input_token_cost": 5.0001e-07, @@ -77877,7 +77889,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "global.moonshotai.kimi-k3": { "supports_regex_lookaround": false, @@ -77907,7 +77921,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "us.moonshotai.kimi-k3": { "supports_regex_lookaround": false, @@ -77937,7 +77953,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "openrouter/xiaomi/mimo-v2.6-flash": { "cache_read_input_token_cost": 2.8e-09, diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index f2c5d066714..5a33fc0b256 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -2,10 +2,12 @@ Unified /v1/messages endpoint - (Anthropic Spec) """ +from collections.abc import Mapping from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse +from pydantic import JsonValue import litellm from litellm.anthropic_interface.exceptions import ( @@ -27,6 +29,7 @@ from litellm.proxy.common_request_processing import ( create_response, log_llm_api_exception, proxy_exception_from_http_exception, + proxy_exception_from_route_error, resolve_litellm_call_id, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id @@ -41,7 +44,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) -from litellm.types.utils import TokenCountResponse +from litellm.proxy.common_utils.token_count_settings import count_source_fields router: Final = APIRouter() @@ -289,6 +292,14 @@ async def anthropic_response( ) +def _requested_output_format(data: Mapping[str, JsonValue]) -> JsonValue: + output_format: Final = data.get("output_format") + output_config: Final = data.get("output_config") + if isinstance(output_format, dict) or not isinstance(output_config, dict): + return output_format + return output_config.get("format") + + @router.post( "/v1/messages/count_tokens", tags=["[beta] Anthropic Messages Token Counting"], @@ -318,7 +329,7 @@ async def count_tokens( Returns: {"input_tokens": } """ - from litellm.proxy.proxy_server import token_counter as internal_token_counter + from litellm.proxy.proxy_server import count_request_tokens litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: @@ -330,10 +341,10 @@ async def count_tokens( messages: Final = data.get("messages", []) if not model_name: - raise HTTPException(status_code=400, detail={"error": "model parameter is required"}) + raise HTTPException(status_code=400, detail="model parameter is required") if not messages: - raise HTTPException(status_code=400, detail={"error": "messages parameter is required"}) + raise HTTPException(status_code=400, detail="messages parameter is required") # Create TokenCountRequest for the internal endpoint from litellm.proxy._types import TokenCountRequest @@ -345,35 +356,21 @@ async def count_tokens( system=data.get("system"), ) - # Call the internal token counter function with direct request flag set to False - token_response: Final = await internal_token_counter( + token_response: Final = await count_request_tokens( request=token_request, call_endpoint=True, + messages_format="anthropic", + output_format=_requested_output_format(data), ) - _token_response_dict: dict = {} - if isinstance(token_response, TokenCountResponse): - _token_response_dict = token_response.model_dump() - elif isinstance(token_response, dict): - _token_response_dict = token_response - - # Convert the internal response to Anthropic API format - return {"input_tokens": _token_response_dict.get("total_tokens", 0)} - - except HTTPException: - raise + except HTTPException as e: + return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) except ProxyException as e: - status_code: Final = int(e.code) if e.code and e.code.isdigit() else 500 - detail: Final = AnthropicExceptionMapping.transform_to_anthropic_error( - status_code=status_code, - raw_message=e.message, - ) - raise HTTPException( - status_code=status_code, - detail=detail, - ) + return _anthropic_error_json_response(e, request) except Exception as e: log_llm_api_exception(e, litellm_call_id) - raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) + return _anthropic_error_json_response(proxy_exception_from_route_error(e), request) + else: + return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)} @router.post( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2b65d713491..14bd302d314 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -637,6 +637,27 @@ def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, s ) +def proxy_exception_from_route_error(exc: Exception) -> ProxyException: + if isinstance(exc, ProxyException): + return exc + if isinstance(exc, HTTPException): + return proxy_exception_from_http_exception(exc, {}) + error_status: Final = error_status_code(exc, status.HTTP_500_INTERNAL_SERVER_ERROR) + message: Final = attribute_of(exc, "message", str(exc)) + carried_code: Final = attribute_of(exc, "code") + provider_fields: Final = attribute_of(exc, "provider_specific_fields") + return ProxyException( + message=redact_internal_details_from_client_message( + strip_bug_report_notice(message) if isinstance(message, str) else str(exc) + ), + type=openai_error_type(exc, error_status), + param=openai_error_param(exc), + code=error_status, + openai_code=carried_code if isinstance(carried_code, str) else None, + provider_specific_fields=provider_fields if isinstance(provider_fields, dict) else None, + ) + + def _collect_response_file_search_vector_store_ids(data: Mapping[str, object]) -> set[str]: vector_store_ids: Final[set[str]] = set() tools: Final = data.get("tools") diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py new file mode 100644 index 00000000000..6f805616820 --- /dev/null +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -0,0 +1,164 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass, replace +from types import MappingProxyType +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, ValidationError + +import litellm +from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.proxy._types import ProxyException +from litellm.types.token_count import TokenCountSourceFields, TokenCountTokenizer +from litellm.types.utils import TokenCountResponse + +CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404}) +TOKENIZER_LOAD_WAIT_SECONDS: Final = 5.0 +TOKENIZER_LOAD_REUSE_SECONDS: Final = 60.0 +TOKENIZER_LOAD_MAX_RUNNING: Final = 4 + + +class _UpstreamErrorDetail(BaseModel): + model_config = ConfigDict(frozen=True, extra="allow") + message: str + + +class _UpstreamErrorEnvelope(BaseModel): + model_config = ConfigDict(frozen=True, extra="allow") + type: Literal["error"] + error: _UpstreamErrorDetail + + +def _joined(message: str, hint: str) -> str: + return f"{message.rstrip().rstrip('.')}. {hint}" + + +def _with_hint(message: str, hint: str) -> str: + try: + envelope: Final = _UpstreamErrorEnvelope.model_validate_json(message) + except ValidationError: + return _joined(AnthropicExceptionMapping.extract_error_message(message), hint) + detail: Final = envelope.error.model_copy(update={"message": _joined(envelope.error.message, hint)}) + return envelope.model_copy(update={"error": detail}).model_dump_json() + + +@dataclass(frozen=True, slots=True) +class TokenCountModel: + model: str + region: str | None + + +@dataclass(frozen=True, slots=True) +class TokenCountSettings: + framing: int | None + count_model: TokenCountModel | None + tokenizer: TokenCountTokenizer | None = None + + +@dataclass(frozen=True, slots=True) +class ProviderCountFailure: + status_code: int + message: str + rejected_before_provider_call: bool = False + + @property + def cannot_count_model(self) -> bool: + return not self.rejected_before_provider_call and self.status_code in CANNOT_COUNT_MODEL_STATUSES + + def with_hint(self, hint: str) -> "ProviderCountFailure": + return replace(self, message=_with_hint(self.message, hint)) + + def to_proxy_exception(self) -> ProxyException: + return ProxyException( + message=self.message, + type="token_counting_error", + param="model", + code=self.status_code, + ) + + +def cost_map_entry(model: str, custom_llm_provider: str | None = None) -> Mapping[str, object]: + try: + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models + return MappingProxyType({}) + + +def _framing(value: object) -> int | None: + match value: + case bool(): + return None + case int() if value >= 0: + return value + case float() if value.is_integer() and value >= 0: + return int(value) + case _: + return None + + +def _non_blank(value: object) -> str | None: + return value if isinstance(value, str) and value.strip() else None + + +def _count_model(value: Mapping[str, object]) -> TokenCountModel | None: + model: Final = _non_blank(value.get("model")) + return None if model is None else TokenCountModel(model=model, region=_non_blank(value.get("region"))) + + +def _geography(region: str) -> str: + return region.rsplit("-", 2)[0] + + +def _within_deployment_geography( + count_model: TokenCountModel | None, deployment_region: Callable[[], str] +) -> TokenCountModel | None: + if count_model is None or count_model.region is None or count_model.model.startswith("global."): + return count_model + if _geography(deployment_region()) == _geography(count_model.region): + return count_model + return replace(count_model, region=None) + + +def _resolved_count_model( + configured: object, default: object, deployment_region: Callable[[], str] +) -> TokenCountModel | None: + configured_mapping: Final = as_str_mapping(configured) + if configured_mapping is not None: + return _count_model(configured_mapping) + default_mapping: Final = as_str_mapping(default) + return ( + None + if default_mapping is None + else _within_deployment_geography(_count_model(default_mapping), deployment_region) + ) + + +def _tokenizer(value: object) -> TokenCountTokenizer | None: + match value: + case "kimi_k3": + return "kimi_k3" + case _: + return None + + +def resolve_token_count_settings( + configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str] +) -> TokenCountSettings: + framing: Final = _framing(configured.get("token_count_framing")) + return TokenCountSettings( + framing=framing if framing is not None else _framing(default.get("token_count_framing")), + count_model=_resolved_count_model( + configured.get("token_count_model"), default.get("token_count_model"), deployment_region + ), + tokenizer=_tokenizer(default.get("token_count_tokenizer")), + ) + + +def count_source_fields(response: TokenCountResponse) -> TokenCountSourceFields: + match response.count_source: + case None: + return {} + case "token_count_model": + return {"litellm_count_source": "token_count_model", "litellm_count_model": response.model_used} + case source: + return {"litellm_count_source": source} diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 8638bc57615..e15daf05ced 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,22 +1,137 @@ +from collections.abc import Mapping +from itertools import chain +from types import MappingProxyType from typing import Final -from fastapi import APIRouter, Depends, Request, Response -from fastapi.responses import ORJSONResponse +from fastapi import APIRouter, Depends, HTTPException, Request, Response +from fastapi.responses import JSONResponse, ORJSONResponse +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types from litellm.proxy._types import * +from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + log_llm_api_exception, + proxy_exception_from_http_exception, + resolve_litellm_call_id, +) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export read_request_body, ) +from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.types.llms.vertex_ai import TokenCountDetailsResponse + +class _GoogleCountTokensTextPart(BaseModel): + model_config = ConfigDict(frozen=True) + text: str + + +class _GoogleCountTokensSystemInstruction(BaseModel): + model_config = ConfigDict(frozen=True) + parts: tuple[_GoogleCountTokensTextPart, ...] + + +class _GoogleCountTokensFunctionDeclaration(BaseModel): + model_config = ConfigDict(frozen=True, populate_by_name=True) + name: str + description: str | None = None + parameters: JsonValue | None = None + parameters_json_schema: JsonValue | None = Field(default=None, alias="parametersJsonSchema") + + +class _GoogleCountTokensTool(BaseModel): + model_config = ConfigDict(frozen=True, populate_by_name=True) + function_declarations: tuple[_GoogleCountTokensFunctionDeclaration, ...] = Field( + default=(), alias="functionDeclarations" + ) + + +class _GoogleCountTokensRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True) + contents: tuple[Mapping[str, JsonValue], ...] = () + system_instruction: _GoogleCountTokensSystemInstruction | None = Field(default=None, alias="systemInstruction") + tools: tuple[_GoogleCountTokensTool, ...] = () + + router: Final = APIRouter( tags=["google genai endpoints"], ) +def _input_schema(declaration: _GoogleCountTokensFunctionDeclaration) -> object: + if declaration.parameters_json_schema is not None: + return declaration.parameters_json_schema + if declaration.parameters is not None: + return normalize_json_schema_types(declaration.parameters) + return {"type": "object", "properties": {}} + + +_GOOGLE_RPC_STATUS: Final = MappingProxyType( + { + 400: "INVALID_ARGUMENT", + 401: "UNAUTHENTICATED", + 403: "PERMISSION_DENIED", + 404: "NOT_FOUND", + 409: "ABORTED", + 429: "RESOURCE_EXHAUSTED", + 499: "CANCELLED", + 501: "UNIMPLEMENTED", + 503: "UNAVAILABLE", + 504: "DEADLINE_EXCEEDED", + } +) + + +def _google_error_response(request: Request, error: ProxyException) -> JSONResponse: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does + ) + + status_code: Final = int(error.code) if error.code is not None and error.code.isdigit() else 500 + _close_dangling_otel_server_span(request, status_code, exc=error) + status: Final = _GOOGLE_RPC_STATUS.get(status_code, "INTERNAL" if status_code >= 500 else "FAILED_PRECONDITION") + return JSONResponse( + status_code=status_code, + content={"error": {"code": status_code, "message": error.message, "status": status}}, + headers=error.headers, + ) + + +def _google_count_tokens_request(model_name: str, data: object) -> TokenCountRequest: + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + validated_data: Final = TypeAdapter(_GoogleCountTokensRequest).validate_python(data) + contents: Final = [dict(content) for content in validated_data.contents] + messages: Final = ( + GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", []) + ) + system: Final = ( + [{"type": "text", "text": part.text} for part in validated_data.system_instruction.parts] + if validated_data.system_instruction is not None + else None + ) + declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools) + tools: Final = [ + { + "name": declaration.name, + **({"description": declaration.description} if declaration.description is not None else {}), + "input_schema": _input_schema(declaration), + } + for declaration in declarations + ] + return TokenCountRequest( + model=model_name, + contents=contents, + messages=messages, + system=system, + tools=tools or None, + ) + + @router.post( "/v1beta/models/{model_name:path}:generateContent", dependencies=[Depends(user_api_key_auth)], @@ -168,36 +283,39 @@ async def google_count_tokens(request: Request, model_name: str): } ``` """ - from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter from litellm.proxy.common_utils.http_parsing_utils import read_request_body from litellm.proxy.proxy_server import token_counter as internal_token_counter data: Final = await read_request_body(request=request) - contents: Final = data.get("contents", []) - # Create TokenCountRequest for the internal endpoint - from litellm.proxy._types import TokenCountRequest + try: + token_request: Final = _google_count_tokens_request(model_name, data) + except ValidationError: + return _google_error_response( + request, + ProxyException( + message="Invalid Google countTokens request body", type="invalid_request_error", param=None, code=400 + ), + ) - # Translate contents to openai format messages using the adapter - messages = GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", []) - - token_request: Final = TokenCountRequest( - model=model_name, - contents=contents, - messages=messages, # compatibility when use openai-like endpoint - ) - - # Call the internal token counter function with direct request flag set to False - token_response: Final = await internal_token_counter( - request=token_request, - call_endpoint=True, - ) - if token_response is None: - return TokenCountDetailsResponse(totalTokens=0, promptTokensDetails=[]) + try: + token_response: Final = await internal_token_counter(request=token_request, call_endpoint=True) + except ProxyException as error: + return _google_error_response(request, error) + except HTTPException as error: + return _google_error_response(request, proxy_exception_from_http_exception(error, {})) + except Exception as error: + log_llm_api_exception(error, resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))) + return _google_error_response( + request, + ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500), + ) original_response: Final[dict] = token_response.original_response or {} - return TokenCountDetailsResponse( - totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0, - promptTokensDetails=original_response.get("promptTokensDetails", []), - ) + counted: Final[TokenCountDetailsResponse] = { + "totalTokens": original_response.get("totalTokens") or token_response.total_tokens or 0, + "promptTokensDetails": original_response.get("promptTokensDetails", []), + **count_source_fields(token_response), + } + return counted # ============================================================ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 67fec38e39e..97001ad7b30 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -47,6 +47,7 @@ from typing import ( ) import anyio +import anyio.lowlevel import websockets import websockets.exceptions from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, ValidationError, with_config @@ -87,6 +88,13 @@ from litellm.litellm_core_utils.litellm_logging import ( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + HUGGING_FACE_SOURCE, + KimiK3Encodings, + kimi_k3_anthropic_tokens, + kimi_k3_template_tokens, + kimi_k3_tokens, +) from litellm.proxy._types import ( UI_TEAM_ID, CallbackDelete, @@ -136,6 +144,15 @@ from litellm.proxy.common_utils.realtime_utils import ( # noqa: F401, RUF100 # _realtime_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export realtime_request_body, ) +from litellm.proxy.common_utils.token_count_settings import ( + TOKENIZER_LOAD_MAX_RUNNING, + TOKENIZER_LOAD_REUSE_SECONDS, + TOKENIZER_LOAD_WAIT_SECONDS, + ProviderCountFailure, + TokenCountSettings, + cost_map_entry, + resolve_token_count_settings, +) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.router_utils.access_windows import access_windows_config_error from litellm.router_utils.add_retry_fallback_headers import ( @@ -165,6 +182,7 @@ from litellm.types.utils import ( ModelResponseStream, StreamingChoices, TextCompletionResponse, + TokenCountMessagesFormat, TokenCountResponse, echoed_cost_map_pricing_fields, is_server_derived_pricing_key, @@ -289,6 +307,7 @@ from litellm.constants import ( CLI_SSO_SESSION_TTL_SECONDS, DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, + DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MODEL_CREATED_AT_TIME, GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_ADMIN_NAME, @@ -318,6 +337,7 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( ) from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, drop_params_flag, get_litellm_metadata_from_kwargs, get_parent_otel_span_from_kwargs, @@ -960,7 +980,13 @@ from litellm.types.secret_managers.main import ( KeyManagementSettings, KeyManagementSystem, ) -from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer, RawRequestTypedDict, StandardLoggingPayload +from litellm.types.utils import ( + CredentialItem, + CustomHuggingfaceTokenizer, + RawRequestTypedDict, + SelectTokenizerResponse, + StandardLoggingPayload, +) from litellm.types.utils import ModelInfo as ModelMapInfo from litellm.utils import add_custom_logger_callback_to_specific_event @@ -14206,52 +14232,367 @@ async def _try_provider_token_count( provider_counter: "BaseTokenCounter", custom_llm_provider: str | None, model_to_use: str, + region: str | None, messages: list | None, contents: list | None, - deployment: dict[str, object] | None, + deployment: Mapping[str, object] | None, request_model: str, tools: list | None = None, system: str | None = None, -) -> Optional["TokenCountResponse"]: - """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.""" + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, +) -> Union["TokenCountResponse", ProviderCountFailure, None]: + """Attempt provider-specific token counting. Returns None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): return None + deployment_in_region: Final = ( + None + if deployment is None + else dict(deployment) + if region is None + else { + **deployment, + "litellm_params": {**(as_str_mapping(deployment.get("litellm_params")) or {}), "aws_region_name": region}, + } + ) try: result: Final = await provider_counter.count_tokens( model_to_use=model_to_use or "", messages=messages, contents=contents, - deployment=deployment, + deployment=deployment_in_region, request_model=request_model, tools=tools, system=system, + messages_format=messages_format, + output_format=output_format, ) except httpx.HTTPStatusError as e: - error_message: Final = getattr(e, "message", None) or str(e) - status_code: Final = getattr(e, "status_code", None) or e.response.status_code - raise ProxyException( - message=error_message, - type="token_counting_error", - param="model", - code=status_code, + return ProviderCountFailure( + status_code=getattr(e, "status_code", None) or e.response.status_code, + message=getattr(e, "message", None) or str(e), ) if result is not None and result.error is True: - if litellm.disable_token_counter is True: - raise ProxyException( - message=result.error_message or "Token counting failed", - type="token_counting_error", - param="model", - code=result.status_code or 500, - ) - verbose_proxy_logger.warning( - "Provider token counting failed (%s): %s. Falling back to local tokenizer.", - result.status_code, - result.error_message, + return ProviderCountFailure( + status_code=result.status_code or 500, + message=result.error_message or "Token counting failed", + rejected_before_provider_call=result.rejected_before_provider_call, ) - return None return result +async def _provider_token_count( + count_with: Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]], + model_to_use: str, + settings: TokenCountSettings, + local_estimate_allowed: bool, +) -> Union["TokenCountResponse", ProviderCountFailure, None]: + count_model: Final = settings.count_model + counted: Final = await ( + count_with(model_to_use, None) if count_model is None else count_with(count_model.model, count_model.region) + ) + if counted is None: + return None + if isinstance(counted, TokenCountResponse): + if count_model is None: + return counted.model_copy(update={"count_source": "provider"}) + return counted.model_copy(update={"count_source": "token_count_model", "model_used": count_model.model}) + failure: Final = ( + counted + if count_model is None or count_model.region is None + else counted.with_hint( + f"This model counts tokens with {count_model.model} in {count_model.region}. " + "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" + ) + ) + if local_estimate_allowed: + return failure + raise failure.to_proxy_exception() + + +@dataclass(frozen=True, slots=True) +class _LocalTokenEstimate: + total_tokens: int + tokenizer_type: str + + +def _is_fork_guard_error(error: BaseException) -> bool: + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking + + return isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)) + + +def _is_tokenizer_panic(error: BaseException) -> bool: + return type(error).__module__ == "pyo3_runtime" and type(error).__name__ == "PanicException" + + +def _answers_provider_error(error: BaseException) -> bool: + return (isinstance(error, Exception) or _is_tokenizer_panic(error)) and not _is_fork_guard_error(error) + + +def _provider_error_after_failed_estimate( + provider_failure: ProviderCountFailure, error: BaseException +) -> ProxyException: + verbose_proxy_logger.warning( + "Local token estimate failed after the provider's count error %s, answering that error: %s", + provider_failure.status_code, + type(error).__name__, + ) + return provider_failure.to_proxy_exception() + + +async def _estimate_or_provider_error(estimate: Awaitable[int], provider_failure: ProviderCountFailure) -> int: + try: + return await estimate + except BaseException as e: + if not _answers_provider_error(e): + raise + raise _provider_error_after_failed_estimate(provider_failure, e) from e + + +@dataclass(frozen=True, slots=True) +class _CustomTokenizerKey: + identifier: str + revision: str + auth_token: str | None + + +class _TokenizerLoadsBusy(Exception): + pass + + +@dataclass(frozen=True, slots=True) +class _TokenizerLoadOutcome: + tokenizer: SelectTokenizerResponse | BaseException + finished_at: float + + +@dataclass(frozen=True, slots=True) +class _SharedTokenizerLoad: + outcome: asyncio.Task[_TokenizerLoadOutcome] + started_at: float + + def reusable(self, now: float) -> bool: + return not self.outcome.done() or now - self.outcome.result().finished_at < TOKENIZER_LOAD_REUSE_SECONDS + + +async def _loaded_tokenizer( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter +) -> SelectTokenizerResponse | BaseException: + try: + return await asyncify(litellm.utils._select_tokenizer, limiter=threads)( + model=model, custom_tokenizer=custom_tokenizer + ) + except asyncio.CancelledError: + raise + except BaseException as e: # noqa: BLE001 # every count sharing this load decides what it answers, pyo3 panics included + return e + + +async def _tokenizer_load_outcome( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter +) -> _TokenizerLoadOutcome: + tokenizer: Final = await _loaded_tokenizer(model, custom_tokenizer, threads) + return _TokenizerLoadOutcome(tokenizer=tokenizer, finished_at=asyncio.get_running_loop().time()) + + +@dataclass(slots=True) +class _TokenizerLoads: + threads: anyio.CapacityLimiter + running: int = 0 + + def start( + self, model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, loop: asyncio.AbstractEventLoop + ) -> _SharedTokenizerLoad | None: + if self.running >= self.threads.total_tokens: + return None + self.running += 1 + outcome: Final = loop.create_task(_tokenizer_load_outcome(model, custom_tokenizer, self.threads)) + outcome.add_done_callback(self._finished) + return _SharedTokenizerLoad(outcome=outcome, started_at=loop.time()) + + def _finished(self, _outcome: asyncio.Task[_TokenizerLoadOutcome]) -> None: + self.running -= 1 + + +_TOKENIZER_LOADS: Final = anyio.lowlevel.RunVar[_TokenizerLoads]("litellm_proxy_tokenizer_loads") + + +def _tokenizer_loads_for_this_loop() -> _TokenizerLoads: + existing: Final = _TOKENIZER_LOADS.get(None) + if existing is not None: + return existing + created: Final = _TokenizerLoads(threads=anyio.CapacityLimiter(TOKENIZER_LOAD_MAX_RUNNING)) + _TOKENIZER_LOADS.set(created) + return created + + +@dataclass(slots=True) +class _TokenizerLoadSlot: + model: str + custom_tokenizer: CustomHuggingfaceTokenizer | None + current: _SharedTokenizerLoad | None = None + + def load(self, loads: _TokenizerLoads, loop: asyncio.AbstractEventLoop) -> _SharedTokenizerLoad | None: + current: Final = self.current + if current is not None and current.reusable(loop.time()): + return current + self.current = loads.start(self.model, self.custom_tokenizer, loop) + return self.current + + +@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) +def _tokenizer_load_slot( + model: str, custom_tokenizer: _CustomTokenizerKey | None, loop: asyncio.AbstractEventLoop +) -> _TokenizerLoadSlot: + return _TokenizerLoadSlot( + model=model, + custom_tokenizer=None + if custom_tokenizer is None + else CustomHuggingfaceTokenizer( + identifier=custom_tokenizer.identifier, + revision=custom_tokenizer.revision, + auth_token=custom_tokenizer.auth_token, + ), + ) + + +async def _tokenizer_after_provider_error( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None +) -> SelectTokenizerResponse | BaseException: + loop: Final = asyncio.get_running_loop() + try: + load: Final = _tokenizer_load_slot( + model, + None + if custom_tokenizer is None + else _CustomTokenizerKey( + identifier=custom_tokenizer["identifier"], + revision=custom_tokenizer["revision"], + auth_token=custom_tokenizer["auth_token"], + ), + loop, + ).load(_tokenizer_loads_for_this_loop(), loop) + if load is None: + return _TokenizerLoadsBusy() + outcome: Final = await asyncio.wait_for( + asyncio.shield(load.outcome), + timeout=max(load.started_at + TOKENIZER_LOAD_WAIT_SECONDS - loop.time(), 0), + ) + except Exception as e: # noqa: BLE001 # a tokenizer that cannot be read, loaded in time, or loaded at all answers the provider's error + return e + return outcome.tokenizer + + +async def _tokenizer_or_provider_error( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, provider_failure: ProviderCountFailure +) -> SelectTokenizerResponse: + tokenizer: Final = await _tokenizer_after_provider_error(model, custom_tokenizer) + if not isinstance(tokenizer, BaseException): + return tokenizer + if not _answers_provider_error(tokenizer): + raise tokenizer.with_traceback(None) + raise _provider_error_after_failed_estimate(provider_failure, tokenizer) from tokenizer + + +async def _local_token_estimate( + model: str, + prompt: str | None, + messages: Sequence[AllMessageValues] | None, + tools: list[ChatCompletionToolParam] | None, + custom_tokenizer: CustomHuggingfaceTokenizer | None, + provider_failure: ProviderCountFailure | None, +) -> _LocalTokenEstimate: + tokenizer: Final = ( + await asyncify(litellm.utils._select_tokenizer)(model=model, custom_tokenizer=custom_tokenizer) + if provider_failure is None + else await _tokenizer_or_provider_error(model, custom_tokenizer, provider_failure) + ) + count: Final = offload_token_count(litellm.token_counter)( + model=model, + text=prompt, + messages=messages, + tools=tools, + custom_tokenizer=tokenizer, + ) + total_tokens: Final = await ( + count if provider_failure is None else _estimate_or_provider_error(count, provider_failure) + ) + return _LocalTokenEstimate(total_tokens=total_tokens, tokenizer_type=str(tokenizer["type"])) + + +_KIMI_K3_ENCODINGS: Final = KimiK3Encodings( + source=HUGGING_FACE_SOURCE, retry_after_seconds=TOKENIZER_LOAD_REUSE_SECONDS +) + + +async def _kimi_k3_count( + request: TokenCountRequest, + messages_format: TokenCountMessagesFormat, + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam] | None, +) -> int | None: + encoding: Final = await _KIMI_K3_ENCODINGS.encoding(TOKENIZER_LOAD_WAIT_SECONDS) + if encoding is None: + return None + match messages_format: + case "anthropic": + return await offload_token_count(kimi_k3_anthropic_tokens)( + encoding, model, request.model_dump(include={"messages", "system", "tools"}, exclude_none=True) + ) + case "openai": + return await offload_token_count(kimi_k3_tokens)(encoding, model, messages, tools) + + +class _TokenCountDeployment(TypedDict, total=False): + litellm_params: ReadOnly[Mapping[str, object]] + model_info: ReadOnly[Mapping[str, object]] + + +_TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment) + + +def _model_group_deployment(router: Router | None, model_group: str) -> Mapping[str, object] | None: + group: Final = () if router is None else router._get_all_deployments(model_name=model_group) + return group[0] if group else None + + +def _non_empty_str(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _deployment_aws_region(deployment: _TokenCountDeployment) -> str: + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + params: Final = deployment.get("litellm_params", {}) + return BaseAWSLLM()._get_aws_region_name( + {"aws_region_name": params.get("aws_region_name")}, model=_non_empty_str(params.get("model")) + ) + + +def _token_count_settings_model(deployment: _TokenCountDeployment, request_model: str) -> tuple[str, str | None]: + model_info: Final = deployment.get("model_info", {}) + params: Final = deployment.get("litellm_params", {}) + base_model: Final = _non_empty_str(model_info.get("base_model")) or _non_empty_str(params.get("base_model")) + if base_model is not None: + return base_model, None + model: Final = _non_empty_str(params.get("model")) + if model is None: + return request_model, None + return model, _non_empty_str(params.get("custom_llm_provider")) + + +_PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( + {litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value} +) + + +def _drops_params(deployment: Mapping[str, object] | None) -> bool: + litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) + return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) + + def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -14274,8 +14615,15 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) Returns: TokenCountResponse """ - global llm_router + return await count_request_tokens(request=request, call_endpoint=call_endpoint) + +async def count_request_tokens( + request: TokenCountRequest, + call_endpoint: bool = False, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, +) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages contents: Final = request.contents @@ -14311,8 +14659,16 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) litellm_model_name = litellm_model_name.split("/", 1)[1] model_to_use: str = ( - litellm_model_name or request.model - ) # use litellm model name, if it's not avalable then fallback to request.model + litellm_model_name if isinstance(litellm_model_name, str) and litellm_model_name else request.model + ) + deployment_view: Final = _TOKEN_COUNT_DEPLOYMENT.validate_python( + deployment or _model_group_deployment(llm_router, request.model) or {} + ) + settings: Final = await asyncify(resolve_token_count_settings)( + configured=deployment_view.get("model_info", {}), + default=await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)), + deployment_region=partial(_deployment_aws_region, deployment_view), + ) # Try provider-specific token counting first - only for non-direct requests (from provider endpoints) provider_counter: BaseTokenCounter | None = None @@ -14323,20 +14679,40 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) if _model is not None: model_to_use = _model - if provider_counter is not None: - result: Final = await _try_provider_token_count( - provider_counter=provider_counter, - custom_llm_provider=custom_llm_provider, - model_to_use=model_to_use, - messages=messages, - contents=contents, - deployment=deployment, - request_model=request.model, - tools=tools, - system=system, + conversation_empty: Final = not messages and not contents + if prompt and not conversation_empty and not _drops_params(deployment): + raise ProxyException( + message="prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt", + type="invalid_request_error", + param="prompt", + code=400, ) - if result is not None: - return result + prompt_only: Final = bool(prompt) and conversation_empty + prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES + ladder: Final = ( + None + if provider_counter is None + else await _provider_token_count( + count_with=partial( + _try_provider_token_count, + provider_counter, + custom_llm_provider, + messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages, + contents=None if prompt_as_message else contents, + deployment=deployment, + request_model=request.model, + tools=tools, + system=system, + messages_format=messages_format, + output_format=output_format, + ), + model_to_use=model_to_use, + settings=settings, + local_estimate_allowed=litellm.disable_token_counter is not True, + ) + ) + if isinstance(ladder, TokenCountResponse): + return ladder # Check if token counter is disabled before fallback if litellm.disable_token_counter is True: @@ -14354,14 +14730,11 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = await asyncify(litellm.utils.select_tokenizer)( - model=model_to_use, custom_tokenizer=custom_tokenizer - ) - - tokenizer_used: Final = str(_tokenizer_used["type"]) system_message: Final = _system_message(system) + prompt_turn: Final = prompt_only and not (system_message is None and not tools) typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes - Sequence[AllMessageValues] | None, messages + Sequence[AllMessageValues] | None, + [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else messages), ) counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages) @@ -14369,18 +14742,38 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None ) - total_tokens: Final = await offload_token_count(litellm.token_counter)( + framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0 + kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None + exact: Final = ( + await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, counted_tools) + if kimi_k3_messages is not None and custom_tokenizer is None + else None + ) + if exact is not None: + return TokenCountResponse( + total_tokens=exact + framing, + request_model=request.model, + model_used=model_to_use, + tokenizer_type="huggingface_tokenizer", + count_source="local_estimate", + ) + estimate: Final = await _local_token_estimate( model=model_to_use, - text=prompt, + prompt=prompt if counted_messages is None else None, messages=counted_messages, tools=counted_tools, - custom_tokenizer=_tokenizer_used, + custom_tokenizer=custom_tokenizer, + provider_failure=ladder, + ) + template: Final = ( + 0 if kimi_k3_messages is None else kimi_k3_template_tokens(len(kimi_k3_messages), bool(counted_tools)) ) return TokenCountResponse( - total_tokens=total_tokens, + total_tokens=estimate.total_tokens + template + framing, request_model=request.model, model_used=model_to_use, - tokenizer_type=tokenizer_used, + tokenizer_type=estimate.tokenizer_type, + count_source="local_estimate", ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 7f9fcca369b..e3bf0f2d693 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg read_request_body, safe_set_request_parsed_body, ) +from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -44,6 +45,7 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ) from litellm.types.responses.main import DeleteResponseResult +from litellm.types.token_count import TokenCountSourceFields from litellm.types.utils import TokenCountResponse if TYPE_CHECKING: @@ -1067,7 +1069,7 @@ class _ResponsesApiErrorBody(TypedDict): error: ReadOnly[_ResponsesApiErrorDetail] -class _ResponsesInputTokensResult(TypedDict): +class _ResponsesInputTokensResult(TokenCountSourceFields): object: ReadOnly[str] input_tokens: ReadOnly[int] @@ -1198,6 +1200,7 @@ async def responses_input_tokens( result: Final[_ResponsesInputTokensResult] = { "object": "response.input_tokens", "input_tokens": token_response.total_tokens, + **count_source_fields(token_response), } return result diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 116fdda0b5d..527da8b2c2b 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -8,6 +8,7 @@ from typing_extensions import ( ) from litellm.types.llms.openai import ChatCompletionFileObject, EmbeddingInput +from litellm.types.token_count import TokenCountSourceFields # Gemini supports nested-list inputs (e.g. [["text", "image"]]) as an explicit # opt-in for combined embeddings — a provider-specific extension of the @@ -319,12 +320,13 @@ class UsageMetadata(TypedDict, total=False): candidatesTokensDetails: list[PromptTokensDetails] # Alternative key name used in some responses -class TokenCountDetailsResponse(TypedDict): +class TokenCountDetailsResponse(TokenCountSourceFields): """ Response structure for token count details with modality breakdown. Example: - {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]} + {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}], + 'litellm_count_source': 'provider'} """ totalTokens: int diff --git a/litellm/types/token_count.py b/litellm/types/token_count.py new file mode 100644 index 00000000000..3b21401066c --- /dev/null +++ b/litellm/types/token_count.py @@ -0,0 +1,16 @@ +from typing import Literal, TypeAlias + +from typing_extensions import NotRequired, ReadOnly, TypedDict + +TokenCountSource: TypeAlias = Literal["provider", "token_count_model", "local_estimate"] +TokenCountTokenizer: TypeAlias = Literal["kimi_k3"] + + +class TokenCountModelInfo(TypedDict): + model: ReadOnly[str] + region: ReadOnly[NotRequired[str]] + + +class TokenCountSourceFields(TypedDict, total=False): + litellm_count_source: ReadOnly[TokenCountSource] + litellm_count_model: ReadOnly[str] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 0fef58ef673..9ced1cd844a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -9,6 +9,7 @@ from typing import ( Any, Final, Literal, + TypeAlias, get_args, ) @@ -56,6 +57,7 @@ from litellm.types.llms.base import ( LiteLLMPydanticObjectBase, ) from litellm.types.mcp import MCPServerCostInfo +from litellm.types.token_count import TokenCountModelInfo, TokenCountSource, TokenCountTokenizer from litellm.types.workload_identity import ( ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS, @@ -227,6 +229,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False): bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None bedrock_converse_supports_strict_tools: bool | None supports_regex_lookaround: ReadOnly[bool | None] + token_count_framing: ReadOnly[int | None] + token_count_model: ReadOnly[TokenCountModelInfo | None] + token_count_tokenizer: ReadOnly[TokenCountTokenizer | None] class SearchContextCostPerQuery(TypedDict, total=False): @@ -3958,11 +3963,14 @@ DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_p DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset({"supports_regex_lookaround"}) +DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset({"token_count_framing", "token_count_model"}) + SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__) - frozenset(CustomPricingLiteLLMParams.model_fields) - DEPLOYMENT_SCOPED_PRICING_FIELDS - DEPLOYMENT_SCOPED_CAPABILITY_FIELDS + - DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS ) @@ -4364,11 +4372,15 @@ class LiteLLMLoggingBaseClass: pass +TokenCountMessagesFormat: TypeAlias = Literal["openai", "anthropic"] + + class TokenCountResponse(LiteLLMPydanticObjectBase): total_tokens: int request_model: str model_used: str tokenizer_type: str + count_source: TokenCountSource | None = None original_response: dict | None = None """ Original Response from upstream API call - if an API call was made for token counting @@ -4379,6 +4391,7 @@ class TokenCountResponse(LiteLLMPydanticObjectBase): HTTP status code from the token counting API (e.g., 200 for success, 429 for rate limit, 400 for bad request) """ status_code: int | None = None + rejected_before_provider_call: bool = Field(default=False, exclude=True) class CustomHuggingfaceTokenizer(TypedDict): diff --git a/litellm/utils.py b/litellm/utils.py index 059f7e5e8de..b01aec72283 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6429,6 +6429,9 @@ def get_model_info_helper( uses_embed_content=_model_info.get("uses_embed_content", None), supports_image_size=_model_info.get("supports_image_size", None), supported_audio_formats=_model_info.get("supported_audio_formats", None), + token_count_framing=_model_info.get("token_count_framing", None), + token_count_model=_model_info.get("token_count_model", None), + token_count_tokenizer=_model_info.get("token_count_tokenizer", None), vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None), ) for cost_key, cost_value in _model_info.items(): diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8d0157f0b97..9de0a3d967f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -20077,7 +20077,13 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "token_count_framing": 97 + }, + "databricks/system.ai.deepseek-v4-flash-0731": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 97 }, "databricks/databricks-deepseek-v4-pro-0813": { "cache_creation_input_token_cost": 1.31999e-06, @@ -21068,7 +21074,13 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 100 + }, + "databricks/system.ai.kimi-k3": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 100 }, "databricks/databricks-llama-4-maverick": { "cache_creation_input_token_cost": 5.0001e-07, @@ -77877,7 +77889,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "global.moonshotai.kimi-k3": { "supports_regex_lookaround": false, @@ -77907,7 +77921,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "us.moonshotai.kimi-k3": { "supports_regex_lookaround": false, @@ -77937,7 +77953,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3" }, "openrouter/xiaomi/mimo-v2.6-flash": { "cache_read_input_token_cost": 2.8e-09, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index b8b998ebae6..712db682a39 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -1219,6 +1219,18 @@ "additionalProperties": false } }, + "token_count_framing": { + "type": "integer", + "minimum": 0, + "description": "Tokens added once per request to a local count estimate, for the provider's chat framing." + }, + "token_count_tokenizer": { + "type": "string", + "description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.", + "enum": [ + "kimi_k3" + ] + }, "tpm": { "type": "integer", "minimum": 0, diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py new file mode 100644 index 00000000000..74581e44310 --- /dev/null +++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py @@ -0,0 +1,215 @@ +import json +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this synthetic request."}] +_PROVIDER_KEY_VARIABLES: Final = ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "AZURE_AI_API_KEY", + "AZURE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_TENANT_ID", + "OPENAI_API_KEY", +) + + +@pytest.fixture(scope="module") +def provider_error_config(tmp_path_factory: pytest.TempPathFactory) -> Path: + config: Final = tmp_path_factory.mktemp("count-tokens-provider-errors") / "config.yaml" + base: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" + config.write_text(base.read_text() + "\nlitellm_settings:\n disable_token_counter: true\n") + return config + + +@pytest.fixture(scope="module") +def keyless_proxy(tmp_path_factory: pytest.TempPathFactory, provider_error_config: Path) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-tokens-error-shapes"), + {}, + config=provider_error_config, + remove_environment=_PROVIDER_KEY_VARIABLES, + workers=1, + ) as owned, + ): + yield owned + + +@pytest.fixture(scope="module") +def azure_environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-tokens-azure-environment-key"), + {"AZURE_API_KEY": "synthetic-azure-env-key"}, + remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name != "AZURE_API_KEY"), + workers=1, + ) as owned, + ): + yield owned + + +def _anthropic(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client + ) + + +def _never_called(request: Request) -> Reply: + raise AssertionError(f"count reached the provider without a key: {request.target}") + + +@pytest.mark.parametrize( + ("deployment", "missing"), + ( + ({"model": "anthropic/claude-sonnet-4-6"}, "No Anthropic API key is set for model "), + ({"model": "azure_ai/claude-sonnet-4-6"}, "No Azure AI API key or Entra ID credential is set for model "), + ), + ids=("anthropic", "azure_ai_anthropic"), +) +def test_anthropic_count_without_a_provider_credential_is_an_anthropic_401( + keyless_proxy: OwnedProxy, deployment: dict[str, JsonValue], missing: str +) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(**deployment, api_base=wire.url, api_key=None, num_retries=0) + with pytest.raises(anthropic.AuthenticationError) as raised: + _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert raised.value.body == { + "type": "error", + "error": {"type": "authentication_error", "message": raised.value.body["error"]["message"]}, + } + assert missing in raised.value.message, raised.value.message + assert wire.drain() == () + + +def test_azure_ai_anthropic_count_sends_the_azure_environment_key_as_x_api_key( + azure_environment_key_proxy: OwnedProxy, +) -> None: + gateway: Final = azure_environment_key_proxy.gateway + + def count(request: Request) -> Reply: + assert request.headers.get("x-api-key") == "synthetic-azure-env-key", request.headers + return Reply(body=json.dumps({"input_tokens": 58}).encode()) + + with wire_server(count) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-sonnet-4-6", + api_base=wire.url, + api_key=None, + num_retries=0, + ) + result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert result.input_tokens == 58 + + +def test_azure_ai_anthropic_count_authenticates_with_an_entra_token(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + + def count(request: Request) -> Reply: + assert request.headers.get("authorization") == "Bearer synthetic-entra-token", request.headers + assert "x-api-key" not in request.headers, request.headers + return Reply(body=json.dumps({"input_tokens": 61}).encode()) + + with wire_server(count) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-sonnet-4-6", + api_base=wire.url, + api_key=None, + azure_ad_token="synthetic-entra-token", + num_retries=0, + ) + result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert result.input_tokens == 61 + assert len(wire.drain()) == 1 + + +def test_anthropic_count_request_errors_use_the_anthropic_envelope(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + upstream_error: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted reject"}} + with ( + wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", api_base=wire.url, api_key="synthetic-key", num_retries=0 + ) + missing_messages: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model}) + assert missing_messages.status_code == 400, missing_messages.text + assert _JSON_OBJECT.validate_json(missing_messages.content) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "messages parameter is required"}, + } + + rejected: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _MESSAGES} + ) + assert rejected.status_code == 400, rejected.text + body: Final = _JSON_OBJECT.validate_json(rejected.content) + assert set(body) <= {"type", "error", "request_id"}, body + assert body["error"] == upstream_error["error"], body + assert len(wire.drain()) == 1 + + +def test_openai_input_tokens_without_a_provider_key_is_an_openai_401(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key=None, num_retries=0 + ) + client: Final = openai.OpenAI( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client + ) + with pytest.raises(openai.AuthenticationError) as raised: + client.responses.input_tokens.count(model=model, input="Count this synthetic request.") + assert f"No OpenAI API key is set for model {model}" in raised.value.message + assert wire.drain() == () + + +def test_google_count_tokens_errors_use_the_google_envelope(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + upstream_error: Final = { + "error": {"code": 400, "message": "API key not valid. Please pass a valid API key.", "status": "INVALID_ARGUMENT"} + } + with ( + wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="gemini/gemini-2.5-flash", api_base=wire.url, api_key="synthetic-key", num_retries=0 + ) + rejected: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + {"contents": [{"role": "user", "parts": [{"text": "Count this synthetic request."}]}]}, + ) + assert rejected.status_code == 400, rejected.text + error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"] + assert isinstance(error, dict) and error["code"] == 400 and error["status"] == "INVALID_ARGUMENT", error + assert "API key not valid" in str(error["message"]), error + assert len(wire.drain()) == 1 + + malformed: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", {"contents": "not-a-list"}) + assert malformed.status_code == 400, malformed.text + assert _JSON_OBJECT.validate_json(malformed.content) == { + "error": {"code": 400, "message": "Invalid Google countTokens request body", "status": "INVALID_ARGUMENT"} + } diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py new file mode 100644 index 00000000000..2b23ccdf4f2 --- /dev/null +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -0,0 +1,393 @@ +import json +from collections.abc import Iterator, Mapping +from typing import Final + +import pytest +from integration._support.client import Scenario, gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_COUNT: Final = 4242 +_TOOL_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {}} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@pytest.fixture(scope="module") +def counting_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process(gateway, tmp_path_factory.mktemp("openai-count-tools"), {}, workers=1) as owned, + ): + yield owned + + +def _deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model( + model="openai/gpt-4o-mini", + api_base=f"{api_base}/v1", + api_key="sk-integration-openai", + ) + + +def _count_response(_request: Request) -> Reply: + return Reply(body=json.dumps({"input_tokens": _COUNT}).encode()) + + +_PDF_BASE64: Final = "JVBERi0xLjQK" +_PNG_BASE64: Final = "aW1hZ2U=" +_LOOKUP_TOOL: Final[dict[str, JsonValue]] = { + "name": "lookup", + "description": "Find a record", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, +} +_RESPONSES_REPLY: Final[Mapping[str, JsonValue]] = { + "id": "resp_count_parity", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_count_parity", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "OK", "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}, +} +_WRAPPED_REASONING: Final = "litellm_enc:bW9kZWxfaWQ6ZGVwbG95bWVudC0x;gAAAAoriginal" +_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}} +_ANTHROPIC_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = { + "structured-output": { + "output_config": {"format": {"type": "json_schema", "schema": _ANSWER_SCHEMA}}, + "messages": [{"role": "user", "content": "Answer in JSON."}], + }, + "affinity-wrapped-reasoning": { + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{_WRAPPED_REASONING}"}, + {"type": "text", "text": "OK"}, + ], + }, + {"role": "user", "content": "again"}, + ], + }, + "base64-pdf": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the attached contract."}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": _PDF_BASE64}, + }, + ], + } + ], + }, + "base64-image": { + "system": [{"type": "text", "text": "Describe images briefly."}], + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": _PNG_BASE64}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + }, + "replayed-web-search": { + "messages": [ + {"role": "user", "content": "What is the newest LiteLLM release?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "newest LiteLLM release"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": "LiteLLM ships a weekly release candidate and a stable line.", + } + ], + }, + {"type": "text", "text": "The newest release is in the release notes."}, + ], + }, + {"role": "user", "content": "Thanks."}, + ], + }, + "tool-turns": { + "system": "Use the lookup tool.", + "tools": [_LOOKUP_TOOL], + "messages": [ + {"role": "user", "content": "Find record 14."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Looking it up."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "14"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}], + }, + ], + }, +} + + +def _responses_peer(request: Request) -> Reply: + if request.target == "/v1/responses/input_tokens": + return _count_response(request) + if request.target == "/v1/responses": + return Reply(body=json.dumps(_RESPONSES_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": request.target}).encode()) + + +def _body_for(wire_requests: tuple[Request, ...], target: str) -> dict[str, JsonValue]: + matching: Final = tuple(request for request in wire_requests if request.target == target) + assert len(matching) == 1, [request.target for request in wire_requests] + return _JSON_OBJECT.validate_json(matching[0].body) + + +def _counted_fields(body: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + return {key: body[key] for key in ("input", "instructions", "tools", "text") if body.get(key) is not None} + + +def _count_body(wire: Wire) -> dict[str, JsonValue]: + requests: Final = wire.drain() + assert len(requests) == 1, requests + request: Final = requests[0] + assert request.method == "POST" + assert request.target == "/v1/responses/input_tokens", request.target + return _JSON_OBJECT.validate_json(request.body) + + +def test_openai_count_tokens_uses_deployment_api_base(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": [{"role": "user", "content": "hello"}]}, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT + assert "tools" not in _count_body(wire) + + +def test_google_count_tokens_maps_anthropic_tools_to_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + { + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "tools": [ + { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + } + ] + } + ], + }, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["totalTokens"] == _COUNT + assert _count_body(wire)["tools"] == [ + { + "type": "function", + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + "strict": False, + } + ] + + +def test_messages_count_tokens_maps_anthropic_tools_to_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + { + "model": model, + "messages": [{"role": "user", "content": "hello"}], + "tools": [ + { + "name": "lookup", + "description": "Find a record", + "input_schema": _TOOL_SCHEMA, + } + ], + }, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT + assert _count_body(wire)["tools"] == [ + { + "type": "function", + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + "strict": False, + } + ] + + +def test_utils_token_counter_keeps_chat_tool_conversion(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/utils/token_counter", + { + "model": model, + "messages": [{"role": "user", "content": "hello"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + }, + } + ], + }, + params={"call_endpoint": "true"}, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["total_tokens"] == _COUNT + assert _count_body(wire)["tools"] == [ + { + "type": "function", + "name": "lookup", + "description": "Find a record", + "parameters": _TOOL_SCHEMA, + } + ] + + +@pytest.mark.parametrize("case", sorted(_ANTHROPIC_BODIES)) +def test_messages_count_tokens_sends_the_input_v1_messages_inference_sends( + counting_proxy: OwnedProxy, case: str +) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + body: Final = {"model": model, **_ANTHROPIC_BODIES[case]} + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT + requests: Final = wire.drain() + inference_body: Final = _body_for(requests, "/v1/responses") + count_body: Final = _body_for(requests, "/v1/responses/input_tokens") + assert _counted_fields(count_body) == _counted_fields(inference_body) + assert count_body["model"] == inference_body["model"] + + +_UNCONVERTIBLE_BODY: Final[Mapping[str, JsonValue]] = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}}, + {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "search_result", + "source": "https://example.com/hits", + "title": "Hits", + "content": [{"type": "text", "text": "hit"}], + } + ], + } + ], + }, + ], +} + + +def test_messages_count_tokens_rejects_blocks_it_cannot_convert(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, **_UNCONVERTIBLE_BODY}) + assert response.status_code == 400, response.text + assert "document with a text source, search_result, search_result inside a tool_result" in response.text + assert wire.drain() == () + + +def test_messages_count_tokens_without_a_role_answers_400_before_calling_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": [{"role": "user", "content": "hello"}, {"content": "no role"}]}, + ) + assert response.status_code == 400, response.text + assert "messages.1.role: Field required" in response.text + assert wire.drain() == () + + +def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + { + "systemInstruction": {"parts": [{"text": "Be brief."}, {"text": "Answer in English."}]}, + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + }, + ) + assert response.status_code == 200, response.text + assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English." diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py new file mode 100644 index 00000000000..5fe82d37e9e --- /dev/null +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -0,0 +1,596 @@ +import asyncio +import json +import threading +from collections.abc import Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from types import MappingProxyType +from typing import Final +from uuid import uuid4 + +import anthropic +import httpx +import pytest +from integration._support.client import Gateway, Scenario +from integration._support.process import graceful_stop_seconds +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_MODEL: Final = "gemini-4-argon" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_COUNT: Final = 777 +_UNIT: Final = " red blue green yellow orange purple black white silver gold" +_USER: Final = f"{_UNIT * 100}\nReturn exactly OK." +_MESSAGES: Final[JsonValue] = [{"role": "user", "content": _USER}] +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) +_REPLY: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "OK"}]}, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 1, "totalTokenCount": 11}, +} + + +class _GeminiPeer: + def __init__(self) -> None: + self.status: int = 200 + self.requests: tuple[Request, ...] = () + self._lock: Final = threading.Lock() + + def set_status(self, status: int) -> None: + self.status = status + + def __call__(self, request: Request) -> Reply: + with self._lock: + self.requests = (*self.requests, request) + status: Final = self.status + if request.target.endswith(":countTokens"): + return Reply(status=status, body=json.dumps({"totalTokens": _COUNT}).encode()) + if request.target.endswith(":generateContent"): + return Reply(body=json.dumps(_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": request.target}).encode()) + + def drain(self) -> tuple[Request, ...]: + with self._lock: + requests: Final = self.requests + self.requests = () + return requests + + +def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str: + return scenario.model( + model=f"vertex_ai/{_MODEL}", + api_base=api_base, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + ) + + +def _messages(model: str, system: JsonValue | None, tools: JsonValue | None) -> Mapping[str, JsonValue]: + return { + "model": model, + "messages": _MESSAGES, + **({"system": system} if system is not None else {}), + **({"tools": tools} if tools is not None else {}), + } + + +def _json_body(request: Request) -> Mapping[str, JsonValue]: + return _JSON_OBJECT.validate_json(request.body) + + +def _count_requests(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if request.target.endswith(":countTokens")) + + +def _count_tokens_tools(generated_tools: JsonValue) -> JsonValue: + return _JSON_LIST.validate_python( + [ + { + { + "function_declarations": "functionDeclarations", + "url_context": "urlContext", + "code_execution": "codeExecution", + }.get(key, key): value + for key, value in tool.items() + } + for tool in TypeAdapter(list[Mapping[str, JsonValue]]).validate_python(generated_tools) + ] + ) + + +def _count_tokens_request_from_generate_content( + generated_body: Mapping[str, JsonValue], +) -> Mapping[str, JsonValue]: + system_instruction: Final = generated_body.get("systemInstruction", generated_body.get("system_instruction")) + generated_tools: Final = generated_body.get("tools") + return MappingProxyType( + { + "contents": generated_body["contents"], + **({"systemInstruction": system_instruction} if system_instruction is not None else {}), + **({"tools": _count_tokens_tools(generated_tools)} if generated_tools is not None else {}), + } + ) + + +def _tools(description: str) -> JsonValue: + return _JSON_LIST.validate_python( + [ + { + "name": "noop", + "description": description, + "input_schema": {"type": "object", "properties": {}}, + } + ] + ) + + +@pytest.mark.parametrize( + ("large_system", "large_tool"), + ((False, False), (True, False), (False, True), (True, True)), + ids=("baseline", "large-system", "large-tool", "large-system-and-tool"), +) +def test_messages_and_count_tokens_send_matching_gemini_context( + gateway: Gateway, large_system: bool, large_tool: bool +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + system: Final = ( + f"{uuid4().hex}\n" + "This is a synthetic context-budget diagnostic. Return exactly OK." + (_UNIT * 100 if large_system else "") + if large_system + else None + ) + tools: Final = _tools(_UNIT * 100) if large_tool else None + body: Final = _messages(model, system, tools) + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + requests: Final = peer.drain() + generated_body: Final = _json_body( + next(request for request in requests if request.target.endswith(":generateContent")) + ) + count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens"))) + assert count_body == _count_tokens_request_from_generate_content(generated_body) + + +_LOOKUP_TOOL: Final[JsonValue] = { + "name": "lookup", + "description": "Look up a value", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, +} +_CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "base64-pdf": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the document."}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + }, + ], + } + ], + }, + "base64-image": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + }, + "tool-turns": { + "tools": [_LOOKUP_TOOL], + "messages": [ + {"role": "user", "content": "Look up the value."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Looking it up."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "value"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}], + }, + ], + }, + } +) + + +def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key") + + +def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: + if "contents" in count_body: + return count_body["contents"] + return _JSON_OBJECT.validate_python(count_body["generateContentRequest"])["contents"] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +@pytest.mark.parametrize("case", sorted(_CONTENT_BLOCK_BODIES)) +def test_count_tokens_sends_the_contents_messages_inference_sends(gateway: Gateway, provider: str, case: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + body: Final = {"model": model, **_CONTENT_BLOCK_BODIES[case]} + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + requests: Final = peer.drain() + generated_body: Final = _json_body( + next(request for request in requests if request.target.endswith(":generateContent")) + ) + count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens"))) + assert _counted_contents(count_body) == generated_body["contents"] + + +_UNCONVERTIBLE_MESSAGES: Final[JsonValue] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}}, + {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "search_result", + "source": "https://example.com/hits", + "title": "Hits", + "content": [{"type": "text", "text": "hit"}], + } + ], + } + ], + }, +] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_count_tokens_rejects_blocks_it_cannot_convert(gateway: Gateway, provider: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _UNCONVERTIBLE_MESSAGES} + ) + assert response.status_code == 400, response.text + assert "document with a text source, search_result, search_result inside a tool_result" in response.text + assert peer.drain() == () + + +def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + system: Final[JsonValue] = [ + {"type": "text", "text": "First instruction"}, + {"type": "text", "text": "Second instruction"}, + ] + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, system, None)) + assert response.status_code == 200, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + assert _json_body(requests[0]) == { + "contents": [{"role": "user", "parts": [{"text": _USER}]}], + "systemInstruction": {"parts": [{"text": "First instruction"}, {"text": "Second instruction"}]}, + } + + +def test_count_tokens_without_system_or_tools_sends_contents_only(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, None)) + assert response.status_code == 200, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + assert _json_body(requests[0]) == {"contents": [{"role": "user", "parts": [{"text": _USER}]}]} + + +_SDK_COUNT_ARGS: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "messages": [{"role": "user", "content": _USER}], + "system": "Follow the system instruction.", + "tools": [ + { + "name": "noop", + "description": "Return no output", + "input_schema": {"type": "object", "properties": {}}, + } + ], + } +) + + +def _sync_sdk_count_tokens(gateway: Gateway, model: str) -> int: + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + return client.messages.count_tokens(model=model, **_SDK_COUNT_ARGS).input_tokens # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping + + +async def _async_sdk_count_tokens(gateway: Gateway, model: str) -> int: + client: Final = anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + return (await client.messages.count_tokens(model=model, **_SDK_COUNT_ARGS)).input_tokens # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync-sdk", "async-sdk")) +def test_anthropic_sdk_count_tokens_forwards_system_and_tools(gateway: Gateway, asynchronous: bool) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + input_tokens: Final = ( + asyncio.run(_async_sdk_count_tokens(gateway, model)) + if asynchronous + else _sync_sdk_count_tokens(gateway, model) + ) + assert input_tokens == _COUNT + requests: Final = _count_requests(wire) + assert len(requests) == 1 + request_body: Final = _json_body(requests[0]) + assert request_body["systemInstruction"] == {"parts": [{"text": "Follow the system instruction."}]} + assert request_body["tools"] == [ + { + "functionDeclarations": [ + {"name": "noop", "description": "Return no output", "parameters": {"type": "object"}} + ] + } + ] + + +def test_google_count_tokens_forwards_system_and_tools_to_vertex(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + body: Final[Mapping[str, JsonValue]] = { + "contents": [ + {"role": "user", "parts": [{"text": _USER}]}, + {"role": "model", "parts": [{"functionCall": {"name": "noop", "args": {"query": "value"}}}]}, + { + "role": "user", + "parts": [{"functionResponse": {"name": "noop", "response": {"result": "ok"}}}], + }, + ], + "systemInstruction": {"parts": [{"text": "Follow this instruction."}]}, + "tools": [ + { + "functionDeclarations": [ + { + "name": "noop", + "description": "Return no output", + "parameters": {"type": "OBJECT", "properties": {}}, + } + ] + } + ], + } + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["totalTokens"] == _COUNT, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + request_body: Final = _json_body(requests[0]) + assert request_body["contents"] == body["contents"] + assert request_body["systemInstruction"] == body["systemInstruction"] + assert request_body["tools"] == [ + { + "functionDeclarations": [ + { + "name": "noop", + "description": "Return no output", + "parameters": {"type": "object"}, + } + ] + } + ] + + +def test_responses_input_tokens_maps_flat_function_tools_to_vertex(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + body: Final[Mapping[str, JsonValue]] = { + "model": model, + "instructions": "Follow the system instruction.", + "input": [ + { + "type": "function_call", + "call_id": "call_lookup", + "name": "lookup", + "arguments": '{"query":"value"}', + }, + { + "type": "function_call_output", + "call_id": "call_lookup", + "output": "result", + }, + ], + "tools": [ + { + "type": "function", + "name": "lookup", + "description": "Look up a value", + "parameters": {"type": "object", "properties": {}}, + } + ], + } + response: Final = gateway.request("POST", "/v1/responses/input_tokens", body) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + request_body: Final = _json_body(requests[0]) + assert request_body["tools"] == [ + { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Look up a value", + "parameters": {"type": "object"}, + } + ] + } + ] + + +def test_utils_token_counter_maps_openai_function_tools_to_vertex(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + tool_schema: Final[JsonValue] = {"type": "object", "properties": {}} + body: Final[Mapping[str, JsonValue]] = { + "model": model, + "messages": _MESSAGES, + "system": "Follow the system instruction.", + "tools": [ + { + "type": "function", + "function": { + "name": "noop", + "description": "Return no output", + "parameters": tool_schema, + }, + } + ], + } + response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "true"}) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["total_tokens"] == _COUNT, response.text + requests: Final = _count_requests(wire) + assert len(requests) == 1 + request_body: Final = _json_body(requests[0]) + assert request_body["systemInstruction"] == {"parts": [{"text": "Follow the system instruction."}]} + assert request_body["tools"] == [ + { + "functionDeclarations": [ + { + "name": "noop", + "description": "Return no output", + "parameters": {"type": "object"}, + } + ] + } + ] + + +def test_count_tokens_surfaces_peer_errors_and_rejects_non_text_system(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + peer.set_status(400) + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, None)) + assert response.status_code == 400, response.text + assert len(_count_requests(wire)) == 1 + + invalid_peer: Final = _GeminiPeer() + with wire_server(invalid_peer) as invalid_wire, gateway.scenario() as scenario: + invalid_model: Final = _deployment(gateway, scenario, invalid_wire.url) + invalid_response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + _messages( + invalid_model, + [{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + None, + ), + ) + assert invalid_response.status_code == 400, invalid_response.text + assert _count_requests(invalid_wire) == () + + +@pytest.mark.parametrize( + "tools", + ( + [{"name": "noop", "input_schema": {"type": "object", "properties": {"query": "not-a-schema"}}}], + [{"name": "noop", "input_schema": "not-a-schema"}], + ), + ids=("malformed-property", "string-schema"), +) +def test_count_tokens_rejects_unconvertible_tools_without_request(gateway: Gateway, tools: JsonValue) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, tools)) + assert response.status_code == 400, response.text + assert "tools" in response.text + assert _count_requests(wire) == () + + +def _count_wave(clients: Sequence[httpx.Client], body: Mapping[str, JsonValue]) -> tuple[tuple[int, JsonValue], ...]: + def count(client: httpx.Client) -> tuple[int, JsonValue]: + response: Final = client.post("/v1/messages/count_tokens", json=dict(body)) + payload: Final = _JSON_OBJECT.validate_json(response.content) + return response.status_code, payload.get("input_tokens") + + with ThreadPoolExecutor(max_workers=len(clients)) as executor: + return tuple(executor.map(count, clients)) + + +def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario, ExitStack() as stack: + model: Final = _deployment(gateway, scenario, wire.url) + clients: Final = tuple( + stack.enter_context( + httpx.Client( + base_url=str(gateway.client.base_url), + headers={"Authorization": f"Bearer {gateway.key}"}, + timeout=30, + trust_env=False, + ) + ) + for _ in range(3) + ) + body: Final = _messages(model, None, None) + assert _count_wave(clients, body) == ((200, _COUNT), (200, _COUNT), (200, _COUNT)) + peer.set_status(503) + assert _count_wave(clients, body) == ((503, None), (503, None), (503, None)) + peer.set_status(200) + assert _count_wave(clients, body) == ((200, _COUNT), (200, _COUNT), (200, _COUNT)) + requests: Final = _count_requests(wire) + assert len(requests) == 9 + assert ( + tuple(_json_body(request) for request in requests) + == ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9 + ) diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index f958984bae5..6869ac12d98 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -24,6 +24,8 @@ model_list: model: anthropic/claude-opus-5-5 api_base: http://127.0.0.1:8191 api_key: synthetic-anthropic-key + model_info: + token_count_framing: 0 - model_name: anthropic/claude-sonnet-4-6 litellm_params: model: anthropic/claude-sonnet-4-6 diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index 1cc03838f8d..b5145ea49d4 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -399,9 +399,11 @@ def test_in_flight_count_tokens_does_not_reserve_key_budget_away_from_a_completi [{"role": "user", "content": prompt}] ] assert count.status_code == 200, count.text - assert count.json() == {"totalTokens": 12, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}]}, ( - count.text - ) + assert count.json() == { + "totalTokens": 12, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}], + "litellm_count_source": "provider", + }, count.text provider_calls: Final = wire.drain() assert [(call.method, call.target) for call in provider_calls] == [ ("POST", "/v1beta/models/gemini-3.8-flash:countTokens") diff --git a/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py b/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py new file mode 100644 index 00000000000..b5301b413ff --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py @@ -0,0 +1,192 @@ +import asyncio +import json +from collections.abc import Mapping +from typing import Final + +import pytest +from pydantic import JsonValue, TypeAdapter + +from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter +from litellm.llms.anthropic.pass_through.convertible_blocks import ( + unconvertible_anthropic_blocks, + unconvertible_blocks_error, +) +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.llms.vertex_ai.count_tokens.context import gemini_count_contents +from litellm.types.llms.anthropic import AllAnthropicPassThroughMessageValues + +_JSON_OBJECTS: Final = TypeAdapter(list[dict[str, JsonValue]]) + +_PDF: Final = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"} +_PNG: Final = {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="} + +_CONVERTED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object]], ...]] = ( + ("user", {"type": "text", "text": "hi"}), + ("user", {"type": "image", "source": _PNG}), + ("user", {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}), + ("user", {"type": "document", "source": _PDF}), + ("user", {"type": "document", "source": {"type": "url", "url": "https://example.com/a.pdf"}}), + ("user", {"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}), + ("assistant", {"type": "text", "text": "hi"}), + ("assistant", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}), + ("assistant", {"type": "thinking", "thinking": "Let me think.", "signature": "sig"}), +) + +_DROPPED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object], str], ...]] = ( + ("user", {"type": "search_result", "source": "https://example.com", "title": "t", "content": []}, "search_result"), + ("user", {"type": "document", "source": {"type": "text", "data": "notes"}}, "document with a text source"), + ("user", {"type": "image", "source": {"type": "file", "file_id": "file_01"}}, "image with a file source"), + ("user", {"type": "image", "source": {"type": "base64", "data": ""}}, "image with an empty base64 source"), + ("user", {"type": "document"}, "document without a source"), + ("user", {"text": "untyped"}, "a block without a type"), + ("user", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}, "tool_use"), + ( + "assistant", + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {}}, + "server_tool_use", + ), +) + + +def _adapter_messages(role: str, block: Mapping[str, object]) -> list[AllAnthropicPassThroughMessageValues]: + return [{"role": role, "content": [dict(block)]}] # pyright: ignore[reportReturnType] # sends block shapes the TypedDicts leave out on purpose + + +def _converted_by_both_adapters(role: str, block: Mapping[str, object]) -> tuple[bool, bool]: + messages: Final = _adapter_messages(role, block) + chat: Final = _JSON_OBJECTS.validate_python( + LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=messages) + ) + responses: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input(messages) + chat_kept: Final = any( + message.get("content") or message.get("tool_calls") or message.get("thinking_blocks") for message in chat + ) + return chat_kept, bool(responses) + + +@pytest.mark.parametrize(("role", "block"), _CONVERTED_BLOCKS) +def test_blocks_both_conversions_keep_are_accepted(role: str, block: Mapping[str, object]) -> None: + assert _converted_by_both_adapters(role, block) == (True, True) + assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == () + + +@pytest.mark.parametrize(("role", "block", "label"), _DROPPED_BLOCKS) +def test_blocks_both_conversions_drop_are_named(role: str, block: Mapping[str, object], label: str) -> None: + assert _converted_by_both_adapters(role, block) == (False, False) + assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == (label,) + + +def test_labels_are_listed_once_in_order_and_string_content_is_accepted() -> None: + messages: Final = [ + {"role": "user", "content": "plain text"}, + {"role": "user", "content": [{"type": "search_result"}, {"type": "container_upload"}]}, + {"role": "assistant", "content": [{"type": "search_result"}]}, + ] + assert unconvertible_anthropic_blocks(messages) == ("container_upload", "search_result") + + +def test_redacted_thinking_is_accepted_like_other_replayed_thinking() -> None: + messages: Final = [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "opaque"}]}] + assert unconvertible_anthropic_blocks(messages) == () + + +def test_messages_that_are_not_json_are_left_to_the_count() -> None: + messages: Final = [{"role": "user", "content": [{"type": "search_result", "source": object()}]}] + assert unconvertible_anthropic_blocks(messages) == () + + +def _tool_turn(inner: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + return ( + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": [dict(inner)]}]}, + ) + + +def _sent_by_each_counter(inner: Mapping[str, object]) -> tuple[str, str]: + messages: Final = _tool_turn(inner) + gemini: Final = asyncio.run(gemini_count_contents(messages, "anthropic", "gemini-2.5-flash", "gemini", {})) + openai: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input( + [dict(message) for message in messages] # pyright: ignore[reportArgumentType] # sends block shapes the TypedDicts leave out on purpose + ) + return json.dumps(gemini, default=str), json.dumps(openai, default=str) + + +_KEPT_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str], ...]] = ( + ({"type": "text", "text": "forty-two"}, "forty-two"), + ({"type": "image", "source": _PNG}, _PNG["data"]), + ({"type": "document", "source": _PDF}, _PDF["data"]), +) + +_DROPPED_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str, str], ...]] = ( + ( + { + "type": "search_result", + "source": "https://example.com", + "title": "t", + "content": [{"type": "text", "text": "search-snippet"}], + }, + "search-snippet", + "search_result inside a tool_result", + ), + ( + {"type": "document", "source": {"type": "text", "data": "text-notes"}}, + "text-notes", + "document with a text source inside a tool_result", + ), + ({"type": "tool_reference", "tool_name": "find_docs"}, "find_docs", "tool_reference inside a tool_result"), +) + + +@pytest.mark.parametrize(("inner", "marker"), _KEPT_INSIDE_TOOL_RESULT) +def test_tool_result_content_both_counters_send_is_accepted(inner: Mapping[str, object], marker: str) -> None: + gemini, openai = _sent_by_each_counter(inner) + + assert marker in gemini + assert marker in openai + assert unconvertible_anthropic_blocks(_tool_turn(inner)) == () + + +@pytest.mark.parametrize(("inner", "marker", "label"), _DROPPED_INSIDE_TOOL_RESULT) +def test_tool_result_content_both_counters_drop_is_named(inner: Mapping[str, object], marker: str, label: str) -> None: + gemini, openai = _sent_by_each_counter(inner) + + assert marker not in gemini + assert marker not in openai + assert unconvertible_anthropic_blocks(_tool_turn(inner)) == (label,) + + +def _replayed_web_search(encrypted_content: str) -> tuple[Mapping[str, object], ...]: + return ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": encrypted_content, + "snippet": "LiteLLM ships a weekly release candidate.", + } + ], + }, + ], + }, + ) + + +def test_a_replayed_web_search_inference_sends_as_text_is_accepted() -> None: + assert unconvertible_blocks_error(_replayed_web_search(""), "anthropic") is None + + +def test_an_encrypted_web_search_inference_cannot_send_is_named() -> None: + error: Final = unconvertible_blocks_error(_replayed_web_search("opaque"), "anthropic") + assert error is not None + assert error.endswith("server_tool_use, web_search_tool_result") diff --git a/tests/unit/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py index 7f717091754..416cb7dcc9b 100644 --- a/tests/unit/llms/anthropic/test_count_tokens_oauth.py +++ b/tests/unit/llms/anthropic/test_count_tokens_oauth.py @@ -352,3 +352,48 @@ def test_count_token_counters_select_their_own_provider( assert counter.should_use_token_counting_api(custom_llm_provider=provider) is True assert counter.should_use_token_counting_api(custom_llm_provider="unknown") is False + + +@pytest.mark.asyncio +async def test_a_global_anthropic_key_counts_like_chat_would_authenticate(monkeypatch): + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "anthropic_key", FAKE_REGULAR_KEY) + + with respx.mock(assert_all_called=True) as router: + route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 13}) + ) + result = await AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}}, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.total_tokens == 13 + assert route.calls.last.request.headers["x-api-key"] == FAKE_REGULAR_KEY + + +@pytest.mark.asyncio +async def test_no_anthropic_key_returns_a_401_count_error(monkeypatch): + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "anthropic_key", None) + monkeypatch.setattr(litellm, "api_key", None) + + result = await AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}}, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 401 + assert result.total_tokens == 0 diff --git a/tests/unit/llms/azure_ai/anthropic/__init__.py b/tests/unit/llms/azure_ai/anthropic/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/__init__.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py new file mode 100644 index 00000000000..4913d2b0d1d --- /dev/null +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -0,0 +1,217 @@ +import json +import operator +from typing import Final + +import httpx +import pytest +import respx +from pydantic import ValidationError + +import litellm +from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAnthropicTokenCounter +from litellm.types.utils import TokenCountResponse + +API_BASE: Final = "https://my-resource.services.ai.azure.com" +COUNT_URL: Final = f"{API_BASE}/anthropic/v1/messages/count_tokens" + + +async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + respx.post(COUNT_URL).mock(return_value=response) + return await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}}, + request_model="azure_ai/claude-sonnet-4-5", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "expected_total"), + [ + ({"input_tokens": 15}, 15), + ({"input_tokens": "15"}, 15), + ({"input_tokens": 15.0}, 15), + ({"input_tokens": True}, 1), + ({"input_tokens": -3}, -3), + ({}, 0), + ({"type": "error", "error": {"message": "429 rate limit"}}, 0), + ({"input_tokens": 7, "context_management": {"original_input_tokens": 9}, "nested": {"a": [1, None]}}, 7), + ], +) +async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload( + monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": expected_total, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, + "original_response": payload, + "error": False, + "error_message": None, + "status_code": None, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rejected_count", + [ + None, + 15.5, + "Request timed out", + "429 Too Many Requests", + [15], + [*range(403), "Request timed out"], + {"value": 429}, + ], +) +async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message( + monkeypatch: pytest.MonkeyPatch, rejected_count: object +) -> None: + with pytest.raises(ValidationError) as rejection: + TokenCountResponse.model_validate( + { + "total_tokens": rejected_count, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + } + ) + + result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count})) + + assert result is not None + assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)] + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": str(rejection.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False], +) +async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup( + monkeypatch: pytest.MonkeyPatch, payload: object +) -> None: + with pytest.raises(AttributeError) as failed_lookup: + operator.attrgetter("get")(payload) + + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": str(failed_lookup.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None: + assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("upstream_status", "upstream_text"), + [(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")], +) +async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text( + monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": upstream_text, + "status_code": upstream_status, + } + ) + + +@pytest.mark.asyncio +async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with pytest.raises(json.JSONDecodeError) as undecodable: + json.loads("not json") + + result: Final = await _count(monkeypatch, httpx.Response(200, text="not json")) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "azure_ai/claude-sonnet-4-5", + "model_used": "claude-sonnet-4-5", + "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": f"CountTokens processing error: {undecodable.value}", + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("set_base", ("env", "global"), ids=("AZURE_API_BASE", "litellm.api_base")) +async def test_count_uses_the_shared_api_base_chat_would_use(monkeypatch: pytest.MonkeyPatch, set_base: str) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("AZURE_AI_API_BASE", raising=False) + monkeypatch.delenv("AZURE_API_BASE", raising=False) + if set_base == "env": + monkeypatch.setenv("AZURE_API_BASE", API_BASE) + else: + monkeypatch.setattr(litellm, "api_base", API_BASE) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key"}}, + request_model="azure_ai/claude-sonnet-4-5", + ) + assert result is not None + assert result.total_tokens == 6 + assert route.called diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py new file mode 100644 index 00000000000..ffddb09cca7 --- /dev/null +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py @@ -0,0 +1,524 @@ +import asyncio +import base64 +import hashlib +import json +import logging +import re +import threading +from collections.abc import Iterable +from itertools import product +from pathlib import Path +from typing import Final + +import httpx +import pytest +import respx +import tiktoken + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + KimiK3Encodings, + KimiK3TokenizerSource, + kimi_k3_anthropic_tokens, + kimi_k3_tokens, + load_kimi_k3_encoding, +) +from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig + +MODEL: Final = "us.moonshotai.kimi-k3" +_CONTROL: Final = re.compile(r"<\|(?:open|close|sep|end_of_msg)\|>") +_WEATHER: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather in °C", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}, "days": {"type": "integer"}}, + "required": ["city"], + }, + }, +} + + +def _call(call_id: str, arguments: dict[str, object]) -> dict[str, object]: + return {"id": call_id, "type": "function", "function": {"name": "get_weather", "arguments": json.dumps(arguments)}} + + +def _written_ranks(path: Path, tokens: Iterable[bytes]) -> KimiK3TokenizerSource: + path.write_bytes(b"\n".join(base64.b64encode(token) + b" %d" % rank for rank, token in enumerate(tokens))) + return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + + +def _single_bytes() -> tuple[bytes, ...]: + return tuple(bytes((value,)) for value in range(256)) + + +@pytest.fixture(scope="module") +def byte_source(tmp_path_factory: pytest.TempPathFactory) -> KimiK3TokenizerSource: + return _written_ranks(tmp_path_factory.mktemp("bytes") / "tiktoken.model", _single_bytes()) + + +@pytest.fixture(scope="module") +def byte_encoding(byte_source: KimiK3TokenizerSource) -> tiktoken.Encoding: + loaded: Final = load_kimi_k3_encoding(byte_source) + assert loaded is not None + return loaded + + +@pytest.fixture(scope="module") +def pair_encoding(tmp_path_factory: pytest.TempPathFactory) -> tiktoken.Encoding: + pairs: Final = (bytes(pair) for pair in product(range(256), repeat=2)) + source: Final = _written_ranks(tmp_path_factory.mktemp("pairs") / "tiktoken.model", (*_single_bytes(), *pairs)) + loaded: Final = load_kimi_k3_encoding(source) + assert loaded is not None + return loaded + + +def _template_count(encoding: tiktoken.Encoding, rendered: str) -> int: + texts: Final = _CONTROL.split(rendered) + return len(_CONTROL.findall(rendered)) + sum(len(encoding.encode_ordinary(text)) for text in texts) + + +_GENERATION_PROMPT: Final = '<|open|>message role="assistant"<|sep|><|open|>think<|sep|>' + + +# Expected renders follow the Kimi K3 chat template in moonshotai/Kimi-K3 at revision f831ab6, which matched +# Bedrock's prompt_tokens for us.moonshotai.kimi-k3 in 19 live calls on 2026-10-16 +@pytest.mark.parametrize( + ("messages", "tools", "rendered"), + [ + pytest.param( + [{"role": "user", "content": "Hello there"}], + None, + '<|open|>message role="user"<|sep|>Hello there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT, + id="user", + ), + pytest.param( + [ + { + "role": "system", + "content": [{"type": "text", "text": "Be terse."}, {"type": "text", "text": "Answer in French."}], + }, + {"role": "user", "content": [{"type": "text", "text": "Hi"}, {"type": "text", "text": " there"}]}, + ], + None, + '<|open|>message role="system"<|sep|>Be terse.Answer in French.<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="user"<|sep|>Hi there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT, + id="text_parts_join", + ), + pytest.param( + [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "And 3+3?"}, + ], + None, + '<|open|>message role="user"<|sep|>What is 2+2?<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>' + "<|open|>response<|sep|>4<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|>" + '<|open|>message role="user"<|sep|>And 3+3?<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT, + id="assistant_turn", + ), + pytest.param( + [ + {"role": "user", "content": "Weather in Paris and Oslo?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [_call("c1", {"city": "Paris", "days": 3}), _call("c2", {"city": "Oslo"})], + }, + {"role": "tool", "tool_call_id": "c2", "content": "Snow"}, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ], + [_WEATHER], + '<|open|>message role="system" type="tool-declare"<|sep|># Tools\n' + "Here are the available tools, described in JSONSchema.\n\n```json\n" + '[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":' + '{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},' + '"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="user"<|sep|>Weather in Paris and Oslo?<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>' + "<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>" + '<|open|>call tool="get_weather" index="1"<|sep|>' + '<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|>' + '<|open|>argument key="days" type="number"<|sep|>3<|close|>argument<|sep|><|close|>call<|sep|>' + '<|open|>call tool="get_weather" index="2"<|sep|>' + '<|open|>argument key="city" type="string"<|sep|>Oslo<|close|>argument<|sep|><|close|>call<|sep|>' + "<|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>" + '<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="tool" tool="get_weather" index="2"<|sep|>Snow<|close|>message<|sep|><|end_of_msg|>' + + _GENERATION_PROMPT, + id="tool_results_out_of_call_order", + ), + pytest.param( + [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + _call( + "c1", {"a&b": 'x"y', "flag": True, "none": None, "list": [1, "zé"], "object": {"k": True}} + ) + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "Rain"}, + ], + [_WEATHER], + '<|open|>message role="system" type="tool-declare"<|sep|># Tools\n' + "Here are the available tools, described in JSONSchema.\n\n```json\n" + '[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":' + '{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},' + '"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="user"<|sep|>Weather?<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>' + "<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>" + '<|open|>call tool="get_weather" index="1"<|sep|>' + '<|open|>argument key="a&b" type="string"<|sep|>x"y<|close|>argument<|sep|>' + '<|open|>argument key="flag" type="boolean"<|sep|>true<|close|>argument<|sep|>' + '<|open|>argument key="none" type="null"<|sep|>null<|close|>argument<|sep|>' + '<|open|>argument key="list" type="array"<|sep|>[1, "zé"]<|close|>argument<|sep|>' + '<|open|>argument key="object" type="object"<|sep|>{"k": true}<|close|>argument<|sep|>' + "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>" + '<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Rain<|close|>message<|sep|><|end_of_msg|>' + + _GENERATION_PROMPT, + id="argument_types_and_escaping", + ), + ], +) +@pytest.mark.parametrize("encoding_fixture", ["byte_encoding", "pair_encoding"]) +def test_count_is_the_kimi_k3_chat_template_render_of_the_bedrock_request( + request: pytest.FixtureRequest, + encoding_fixture: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + rendered: str, +): + encoding: Final[tiktoken.Encoding] = request.getfixturevalue(encoding_fixture) + + assert kimi_k3_tokens(encoding, MODEL, messages, tools) == _template_count(encoding, rendered) + + +def _parallel_calls(count: int) -> list[dict[str, object]]: + ids: Final = tuple(f"c{number}" for number in range(1, count + 1)) + return [ + {"role": "user", "content": "Weather?"}, + {"role": "assistant", "content": None, "tool_calls": [_call(call_id, {"city": "Paris"}) for call_id in ids]}, + *({"role": "tool", "tool_call_id": call_id, "content": "Sun"} for call_id in ids), + ] + + +def test_tenth_parallel_call_and_its_result_are_numbered_ten(byte_encoding: tiktoken.Encoding): + nine: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(9), [_WEATHER]) + ten: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(10), [_WEATHER]) + assert nine is not None and ten is not None + + assert ten - nine == _template_count( + byte_encoding, + '<|open|>call tool="get_weather" index="10"<|sep|>' + '<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|><|close|>call<|sep|>' + '<|open|>message role="tool" tool="get_weather" index="10"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>', + ) + + +@pytest.mark.parametrize( + "text", + [ + pytest.param("The quick brown fox, jumps over 12345 lazy dogs. " * 60, id="prose"), + pytest.param("a" + "1234567890" * 250 + "1", id="long_digit_run"), + pytest.param("--- " * 600, id="symbols"), + pytest.param("def f(x):\n return x # it's 42\n\n" * 80, id="code"), + pytest.param("word next\t\ttab\n \n" * 120, id="whitespace"), + pytest.param("naïve café, “quotes” & 中文汉字,日本語のテキスト 😀🎉\r\n" * 40, id="unicode"), + pytest.param("We're sure they'd say it's DON'T or I'M " * 60, id="contractions"), + pytest.param("abc123def4567ghi89 " * 120, id="letters_and_digits"), + ], +) +def test_long_text_counts_what_the_tokenizer_counts_for_the_whole_text(pair_encoding: tiktoken.Encoding, text: str): + def counted(content: str) -> int | None: + return kimi_k3_tokens(pair_encoding, MODEL, [{"role": "user", "content": content}], None) + + short: Final = counted("x") + assert short is not None + + assert counted(text) == short - len(pair_encoding.encode_ordinary("x")) + len(pair_encoding.encode_ordinary(text)) + + +@pytest.mark.parametrize( + ("messages", "tools"), + [ + pytest.param( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, + ], + } + ], + None, + id="image", + ), + pytest.param( + [{"role": "user", "content": "Weather?"}, {"role": "tool", "tool_call_id": "c1", "content": "Sun"}], + [_WEATHER], + id="tool_result_without_its_call", + ), + ], +) +def test_request_the_template_cannot_render_is_not_counted( + pair_encoding: tiktoken.Encoding, messages: list[dict[str, object]], tools: list[dict[str, object]] | None +): + assert kimi_k3_tokens(pair_encoding, MODEL, messages, tools) is None + + +@pytest.mark.parametrize( + "tool", + [ + pytest.param( + { + "type": "function", + "name": "get_weather", + "description": "Get the weather in °C", + "parameters": _WEATHER["function"]["parameters"], + }, + id="responses_api", + ), + pytest.param( + { + "name": "get_weather", + "description": "Get the weather in °C", + "input_schema": _WEATHER["function"]["parameters"], + }, + id="anthropic_messages", + ), + ], +) +def test_tool_in_another_api_shape_counts_like_the_same_chat_completions_tool( + pair_encoding: tiktoken.Encoding, tool: dict[str, object] +): + messages: Final = [{"role": "user", "content": "Weather in Paris?"}] + + counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [tool]) + + assert counted is not None + assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER]) + + +@pytest.mark.parametrize( + "tool", + [ + pytest.param({"type": "web_search_20250305", "name": "web_search"}, id="anthropic_web_search"), + pytest.param({"type": "web_search"}, id="responses_web_search"), + pytest.param( + { + "type": "mcp", + "server_label": "docs", + "server_url": "https://mcp.example/mcp", + "require_approval": "never", + }, + id="mcp", + ), + ], +) +def test_tool_bedrock_inference_drops_counts_like_the_function_tools_alone( + pair_encoding: tiktoken.Encoding, tool: dict[str, object] +): + messages: Final = [{"role": "user", "content": "Weather in Paris?"}] + + counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, tool]) + + assert counted is not None + assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER]) + + +def test_responses_custom_tool_counts_like_the_function_tool_inference_sends(pair_encoding: tiktoken.Encoding): + messages: Final = [{"role": "user", "content": "Weather in Paris?"}] + custom_tool: Final = {"type": "custom", "name": "grammar", "description": "Free-form SQL"} + sent_to_inference, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + [dict(custom_tool)] + ) + + without_it: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER]) + + counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, custom_tool]) + + assert counted is not None and without_it is not None + assert counted > without_it + assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, *sent_to_inference]) + + +def test_counting_a_request_with_an_image_url_never_downloads_the_image( + pair_encoding: tiktoken.Encoding, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "user_url_validation", False) + image: Final = respx_mock.get("https://images.example/cat.png").mock(return_value=httpx.Response(200)) + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + {"type": "image_url", "image_url": {"url": "https://images.example/cat.png"}}, + ], + } + ] + + assert kimi_k3_tokens(pair_encoding, MODEL, messages, None) is None + assert not image.called + + +def test_anthropic_request_counts_like_the_same_openai_request(pair_encoding: tiktoken.Encoding): + anthropic: Final = kimi_k3_anthropic_tokens( + pair_encoding, + MODEL, + { + "system": "Be terse.", + "messages": [ + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "c1", "name": "get_weather", "input": {"city": "Paris"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "Sun"}]}, + ], + "tools": [ + { + "name": "get_weather", + "description": "Get the weather in °C", + "input_schema": _WEATHER["function"]["parameters"], + } + ], + }, + ) + openai: Final = kimi_k3_tokens( + pair_encoding, + MODEL, + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": None, "tool_calls": [_call("c1", {"city": "Paris"})]}, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ], + [_WEATHER], + ) + + assert openai is not None + assert anthropic == openai + + +@pytest.mark.parametrize("failure", ["checksum", "missing_file"]) +def test_tokenizer_that_fails_its_checksum_or_cannot_be_read_is_not_loaded( + byte_source: KimiK3TokenizerSource, tmp_path: Path, caplog: pytest.LogCaptureFixture, failure: str +): + source: Final = ( + KimiK3TokenizerSource(ranks_path=byte_source.ranks_path, sha256=hashlib.sha256(b"other").hexdigest()) + if failure == "checksum" + else KimiK3TokenizerSource(ranks_path=lambda: str(tmp_path / "missing"), sha256=byte_source.sha256) + ) + verbose_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + loaded: Final = load_kimi_k3_encoding(source) + finally: + verbose_logger.removeHandler(caplog.handler) + + assert loaded is None + assert any("Kimi K3 tokenizer unavailable" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_slow_tokenizer_load_is_not_waited_on_past_the_wait_and_is_shared_once_done( + byte_source: KimiK3TokenizerSource, +): + released: Final = threading.Event() + + def slow_ranks_path() -> str: + released.wait(10) + return byte_source.ranks_path() + + encodings: Final = KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=slow_ranks_path, sha256=byte_source.sha256), + retry_after_seconds=60, + clock=lambda: 0.0, + ) + + assert await encodings.encoding(wait_seconds=0) is None + released.set() + loaded: Final = await encodings.encoding(wait_seconds=5) + assert loaded is not None + assert await encodings.encoding(wait_seconds=0) is loaded + + +class _Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +async def _loaded_eventually(encodings: KimiK3Encodings) -> tiktoken.Encoding | None: + for _ in range(500): + if (loaded := await encodings.encoding(wait_seconds=0)) is not None: + return loaded + await asyncio.sleep(0.01) + return None + + +@pytest.mark.asyncio +async def test_failed_tokenizer_load_is_retried_in_the_background_only_after_the_retry_window( + byte_source: KimiK3TokenizerSource, tmp_path: Path +): + clock: Final = _Clock() + reads: Final[list[float]] = [] + released: Final = threading.Event() + + def ranks_path() -> str: + reads.append(clock.now) + if len(reads) == 1: + return str(tmp_path / "missing") + released.wait(10) + return byte_source.ranks_path() + + encodings: Final = KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=ranks_path, sha256=byte_source.sha256), + retry_after_seconds=60, + clock=clock, + ) + + assert await encodings.encoding(wait_seconds=5) is None + clock.now = 59.0 + assert await encodings.encoding(wait_seconds=5) is None + clock.now = 60.0 + assert await asyncio.wait_for(encodings.encoding(wait_seconds=5), timeout=1) is None + released.set() + assert await _loaded_eventually(encodings) is not None + assert reads == [0.0, 60.0] + + +@pytest.mark.asyncio +async def test_tokenizer_load_still_running_after_the_wait_is_logged( + byte_source: KimiK3TokenizerSource, caplog: pytest.LogCaptureFixture +): + released: Final = threading.Event() + + def hung_ranks_path() -> str: + released.wait(10) + return byte_source.ranks_path() + + encodings: Final = KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=hung_ranks_path, sha256=byte_source.sha256), retry_after_seconds=60 + ) + verbose_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + loaded: Final = await encodings.encoding(wait_seconds=0.05) + finally: + verbose_logger.removeHandler(caplog.handler) + released.set() + + assert loaded is None + assert any("Kimi K3 tokenizer still loading" in record.getMessage() for record in caplog.records) diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 70946e10590..e54023c6596 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -1,8 +1,200 @@ +from typing import Final from unittest.mock import AsyncMock, patch +import httpx import pytest +from pydantic import JsonValue, TypeAdapter +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.common_utils import GeminiModelInfo, GoogleAIStudioTokenCounter +from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler +from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter as GeminiCountTokensHandler + +_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue]) + + +class _GeminiCountTokensPeer: + def __init__(self) -> None: + self.request_body: dict[str, JsonValue] | None = None + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.request_body = _JSON_BODY.validate_json(request.content) + return httpx.Response(200, json={"totalTokens": 11}, request=request) + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_forwards_system_and_tools() -> None: + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + deployment={ + "litellm_params": { + "api_key": "test-key", + "api_base": "https://count-tokens.test", + "system_instruction": {"parts": [{"text": "deployment-level"}]}, + } + }, + request_model="gemini/gemini-2.5-flash", + tools=[ + { + "name": "lookup", + "description": "Find a record", + "input_schema": {"type": "object", "properties": {}}, + } + ], + system="You are terse.", + ) + + assert result is not None + assert result.error is False + assert result.total_tokens == 11 + assert peer.request_body == { + "generateContentRequest": { + "model": "models/gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "systemInstruction": {"parts": [{"text": "You are terse."}]}, + "tools": [ + { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Find a record", + "parameters": {"type": "object"}, + } + ] + } + ], + } + } + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_rejects_image_system_without_request() -> None: + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + system=[ + {"type": "text", "text": "You are terse."}, + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}, + }, + ], + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert peer.request_body is None + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_rejects_an_image_url_it_refuses_to_fetch_without_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}], + } + ], + contents=None, + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "http://127.0.0.1:9/image.png" in (result.error_message or "") + assert peer.request_body is None + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_answers_with_the_status_google_rejected_it_with() -> None: + rejection: Final = '{"error": {"code": 400, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}' + client: Final = AsyncHTTPHandler( + transport=httpx.MockTransport(lambda request: httpx.Response(400, text=rejection, request=request)) + ) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "Invalid function name" in (result.error_message or "") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("system_instruction", "tools", "expected_body"), + ( + ( + {"parts": [{"text": "Follow the system rule."}]}, + [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "OBJECT"}}]}], + { + "generateContentRequest": { + "model": "models/gemini-4-argon", + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + "systemInstruction": {"parts": [{"text": "Follow the system rule."}]}, + "tools": [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "OBJECT"}}]}], + } + }, + ), + ( + None, + [], + {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}, + ), + ), +) +async def test_gemini_count_tokens_sends_optional_request_fields( + system_instruction: dict[str, JsonValue] | None, + tools: list[dict[str, JsonValue]], + expected_body: dict[str, JsonValue], +) -> None: + peer: Final = _GeminiCountTokensPeer() + token_counter: Final = GeminiCountTokensHandler(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer))) + + result: Final = await token_counter.acount_tokens( + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + model="gemini-4-argon", + api_key="test-key", + api_base="https://count-tokens.test", + system_instruction=system_instruction, + tools=tools, + ) + + assert result == {"totalTokens": 11} + assert peer.request_body == expected_body class TestGeminiModelInfo: @@ -36,18 +228,10 @@ class TestGeminiModelInfo: # Test edge cases where model names end with characters from "models/" # These would be incorrectly processed if using strip("models/") instead of replace("models/", "") models = [ - { - "name": "models/gemini-1.5-pro" - }, # ends with 'o' - would become "gemini-1.5-pr" with strip() - { - "name": "models/test-model" - }, # ends with 'l' - would become "gemini/test-mode" with strip() - { - "name": "models/custom-models" - }, # ends with 's' - would become "gemini/custom-model" with strip() - { - "name": "models/demo" - }, # ends with 'o' - would become "gemini/dem" with strip() + {"name": "models/gemini-1.5-pro"}, # ends with 'o' - would become "gemini-1.5-pr" with strip() + {"name": "models/test-model"}, # ends with 'l' - would become "gemini/test-mode" with strip() + {"name": "models/custom-models"}, # ends with 's' - would become "gemini/custom-model" with strip() + {"name": "models/demo"}, # ends with 'o' - would become "gemini/dem" with strip() ] result = gemini_model_info.process_model_name(models) @@ -98,16 +282,10 @@ class TestGoogleAIStudioTokenCounter: token_counter = GoogleAIStudioTokenCounter() # Test with gemini provider - should return True - assert ( - token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) - is True - ) + assert token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) is True # Test with other providers - should return False - assert ( - token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) - is False - ) + assert token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) is False assert token_counter.should_use_token_counting_api("anthropic") is False assert token_counter.should_use_token_counting_api("vertex_ai") is False @@ -158,7 +336,10 @@ class TestGoogleAIStudioTokenCounter: # Verify the mock was called correctly mock_acount_tokens.assert_called_once_with( - model=model_to_use, contents=contents + model=model_to_use, + contents=contents, + system_instruction=None, + tools=None, ) def test_clean_contents_for_gemini_api_removes_id_field(self): @@ -176,9 +357,7 @@ class TestGoogleAIStudioTokenCounter: "functionResponse": { "id": "read_many_files-1757526647518-730a691aac11c", # This should be removed "name": "read_many_files", - "response": { - "output": "No files matching the criteria were found or all were skipped." - }, + "response": {"output": "No files matching the criteria were found or all were skipped."}, } } ], @@ -187,9 +366,7 @@ class TestGoogleAIStudioTokenCounter: ] # Clean the contents - cleaned_contents = token_counter._clean_contents_for_gemini_api( - contents_with_id - ) + cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_with_id) # Verify the 'id' field was removed function_response = cleaned_contents[1]["parts"][0]["functionResponse"] @@ -198,8 +375,7 @@ class TestGoogleAIStudioTokenCounter: assert "response" in function_response assert function_response["name"] == "read_many_files" assert ( - function_response["response"]["output"] - == "No files matching the criteria were found or all were skipped." + function_response["response"]["output"] == "No files matching the criteria were found or all were skipped." ) def test_clean_contents_for_gemini_api_preserves_other_fields(self): @@ -215,9 +391,7 @@ class TestGoogleAIStudioTokenCounter: ] # Clean the contents - cleaned_contents = token_counter._clean_contents_for_gemini_api( - contents_without_function_response - ) + cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_without_function_response) # Verify the contents are unchanged assert cleaned_contents == contents_without_function_response diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py new file mode 100644 index 00000000000..a45773b93fc --- /dev/null +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -0,0 +1,436 @@ +import json +import operator +from typing import Final + +import httpx +import pytest +import respx +from pydantic import JsonValue, ValidationError + +import litellm +from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse + +COUNT_URL: Final = "https://api.openai.com/v1/responses/input_tokens" + + +async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + respx.post(COUNT_URL).mock(return_value=response) + return await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "expected_total"), + [ + ({"input_tokens": 15}, 15), + ({"input_tokens": "15"}, 15), + ({"input_tokens": 15.0}, 15), + ({"input_tokens": True}, 1), + ({"input_tokens": -3}, -3), + ({}, 0), + ({"object": "response.input_tokens", "error": {"message": "429 rate limit"}}, 0), + ({"input_tokens": 7, "object": "response.input_tokens", "nested": {"a": [1, None]}}, 7), + ], +) +async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload( + monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": expected_total, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "count_source": None, + "original_response": payload, + "error": False, + "error_message": None, + "status_code": None, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rejected_count", + [ + None, + 15.5, + "Request timed out", + "429 Too Many Requests", + [15], + [*range(403), "Request timed out"], + {"value": 429}, + ], +) +async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message( + monkeypatch: pytest.MonkeyPatch, rejected_count: object +) -> None: + with pytest.raises(ValidationError) as rejection: + TokenCountResponse.model_validate( + { + "total_tokens": rejected_count, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + } + ) + + result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count})) + + assert result is not None + assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)] + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": str(rejection.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False], +) +async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup( + monkeypatch: pytest.MonkeyPatch, payload: object +) -> None: + with pytest.raises(AttributeError) as failed_lookup: + operator.attrgetter("get")(payload) + + result: Final = await _count(monkeypatch, httpx.Response(200, json=payload)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": str(failed_lookup.value), + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None: + assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("upstream_status", "upstream_text"), + [(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")], +) +async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text( + monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str +) -> None: + result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text)) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": upstream_text, + "status_code": upstream_status, + } + ) + + +async def _sent_count_body( + monkeypatch: pytest.MonkeyPatch, + messages: list[dict[str, object]], + system: object, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue = None, +) -> object: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1})) + await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=messages, + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + system=system, + messages_format=messages_format, + output_format=output_format, + ) + return json.loads(route.calls.last.request.content) + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_the_document_image_and_tool_turns_to_the_count_endpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + messages: Final[list[dict[str, object]]] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize."}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBE"}}, + {"type": "image", "source": {"type": "url", "url": "https://example.com/page.png"}}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "found"}]}, + ] + + sent: Final = await _sent_count_body( + monkeypatch, messages, [{"type": "text", "text": "Be brief."}], messages_format="anthropic" + ) + + assert sent == { + "model": "gpt-4o", + "input": [ + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize."}, + {"type": "input_file", "filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBE"}, + {"type": "input_image", "image_url": "https://example.com/page.png"}, + ], + }, + {"type": "function_call", "call_id": "toolu_01", "name": "lookup", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "toolu_01", "output": "found"}, + ], + "instructions": "Be brief.", + } + + +@pytest.mark.asyncio +async def test_openai_format_keeps_the_chat_conversion(monkeypatch: pytest.MonkeyPatch) -> None: + messages: Final[list[dict[str, object]]] = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": [{"type": "text", "text": "Summarize."}]}, + {"role": "tool", "tool_call_id": "call_1", "content": "found"}, + ] + + sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="openai") + + assert sent == { + "model": "gpt-4o", + "input": [ + {"role": "user", "content": "Summarize."}, + {"type": "function_call_output", "call_id": "call_1", "output": "found"}, + ], + "instructions": "Be brief.", + } + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_affinity_wrapped_reasoning_unwrapped_as_inference_does( + monkeypatch: pytest.MonkeyPatch, +) -> None: + wrapped: Final = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAoriginal", "deployment-1") # pyright: ignore[reportPrivateUsage] # the wrap inference applies + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"}, + {"type": "text", "text": "OK"}, + ], + }, + {"role": "user", "content": "again"}, + ] + + sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic") + + assert sent == { + "model": "gpt-4o", + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + {"type": "reasoning", "summary": [], "encrypted_content": "gAAAAoriginal"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "OK"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "again"}]}, + ], + } + + +@pytest.mark.asyncio +async def test_anthropic_format_counts_the_structured_output_schema(monkeypatch: pytest.MonkeyPatch) -> None: + schema: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}} + + sent: Final = await _sent_count_body( + monkeypatch, + [{"role": "user", "content": "Answer in JSON."}], + system=None, + messages_format="anthropic", + output_format={"type": "json_schema", "schema": schema}, + ) + + assert sent == { + "model": "gpt-4o", + "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Answer in JSON."}]}], + "text": {"format": {"type": "json_schema", "name": "structured_output", "schema": schema, "strict": False}}, + } + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_a_replayed_web_search_result_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + {"role": "user", "content": "Thanks."}, + ] + + sent: Final = json.dumps(await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic")) + + assert snippet in sent + assert "https://docs.litellm.ai/release_notes" in sent + + +@pytest.mark.asyncio +async def test_anthropic_format_message_without_a_role_answers_400_without_calling_openai( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1})) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}, {"content": "no role"}], + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + messages_format="anthropic", + ) + + assert route.called is False + assert result is not None + assert (result.error, result.status_code, result.error_message) == (True, 400, "messages.1.role: Field required") + + +@pytest.mark.asyncio +async def test_openai_format_sends_a_list_system_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None: + sent: Final = await _sent_count_body( + monkeypatch, + [{"role": "user", "content": "hi"}], + system=[{"type": "text", "text": "Be brief."}, {"type": "text", "text": "Answer in English."}], + messages_format="openai", + ) + + assert sent == { + "model": "gpt-4o", + "input": [{"role": "user", "content": "hi"}], + "instructions": "Be brief.\nAnswer in English.", + } + + +@pytest.mark.asyncio +async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with pytest.raises(json.JSONDecodeError) as undecodable: + json.loads("not json") + + result: Final = await _count(monkeypatch, httpx.Response(200, text="not json")) + + assert result is not None + assert repr(result.model_dump()) == repr( + { + "total_tokens": 0, + "request_model": "openai/gpt-4o", + "model_used": "gpt-4o", + "tokenizer_type": "openai_api", + "count_source": None, + "original_response": None, + "error": True, + "error_message": f"CountTokens processing error: {undecodable.value}", + "status_code": 500, + } + ) + + +@pytest.mark.asyncio +async def test_count_uses_the_global_openai_key_chat_would_use(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "openai_key", "sk-from-litellm-settings") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 9})) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.total_tokens == 9 + assert route.calls.last.request.headers["authorization"] == "Bearer sk-from-litellm-settings" + + +@pytest.mark.asyncio +async def test_count_sends_the_global_key_only_to_the_global_api_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "openai_key", "sk-private-gateway") + monkeypatch.setattr(litellm, "api_base", "https://gateway.internal.example/v1") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with respx.mock: + route: Final = respx.post("https://gateway.internal.example/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 4}) + ) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.total_tokens == 4 + assert route.calls.last.request.headers["authorization"] == "Bearer sk-private-gateway" diff --git a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py index c2efc1acdb9..8850fd13d10 100644 --- a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -1,6 +1,10 @@ +from typing import Final import pytest +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ) @@ -446,3 +450,108 @@ def test_get_required_headers(): assert headers["Authorization"] == "Bearer sk-test-key" assert headers["Content-Type"] == "application/json" + + +PDF_BASE64: Final = "JVBERi0xLjcK" +PNG_BASE64: Final = "iVBORw0KGgo=" + +ANTHROPIC_AGENT_TURN: Final[list[dict[str, object]]] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the indemnity terms."}, + { + "type": "document", + "title": "contract.pdf", + "source": {"type": "base64", "media_type": "application/pdf", "data": PDF_BASE64}, + }, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG_BASE64}}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me pull clause 14."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup_clause", "input": {"id": "14"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}], + }, +] + +ANTHROPIC_TOOLS: Final[list[dict[str, object]]] = [ + { + "name": "lookup_clause", + "description": "Fetch a clause by id", + "input_schema": {"type": "object", "properties": {"id": {"type": "string"}}}, + } +] + + +@pytest.mark.parametrize( + "system", + [None, "You review contracts.", [{"type": "text", "text": "You review contracts."}]], +) +def test_anthropic_messages_convert_exactly_as_v1_messages_inference_converts_them(system: object) -> None: + inference_request: Final[dict[str, object]] = { + "model": "gpt-4o", + "messages": ANTHROPIC_AGENT_TURN, + "tools": ANTHROPIC_TOOLS, + **({"system": system} if system is not None else {}), + } + inference: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_request( + inference_request, include_encrypted_reasoning=False + ) + + counted: Final = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=system, tools=ANTHROPIC_TOOLS + ) + + assert counted == ( + tuple(inference["input"]), + inference.get("instructions"), + tuple(inference["tools"]), + inference.get("text"), + ) + + +def test_anthropic_messages_keep_every_block_the_chat_converter_drops() -> None: + input_items, _, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=None, tools=None + ) + + assert input_items == ( + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize the indemnity terms."}, + { + "type": "input_file", + "filename": "contract.pdf", + "file_data": f"data:application/pdf;base64,{PDF_BASE64}", + }, + {"type": "input_image", "image_url": f"data:image/png;base64,{PNG_BASE64}"}, + ], + }, + {"type": "function_call", "call_id": "toolu_01", "name": "lookup_clause", "arguments": '{"id": "14"}'}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Let me pull clause 14."}], + }, + {"type": "function_call_output", "call_id": "toolu_01", "output": "14. Indemnification."}, + ) + + +def test_anthropic_list_system_becomes_joined_instruction_text() -> None: + _, instructions, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + system=[{"type": "text", "text": "First rule."}, {"type": "text", "text": "Second rule."}], + tools=None, + ) + + assert instructions == "First rule.\nSecond rule." diff --git a/tests/unit/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py index 7d78a2c5238..74cbd7c1144 100644 --- a/tests/unit/llms/openai/test_openai_workload_identity.py +++ b/tests/unit/llms/openai/test_openai_workload_identity.py @@ -630,3 +630,56 @@ class TestClientsideBaseOverride: ) is None ) + + +class TestCountTokens: + @pytest.mark.asyncio + @respx.mock + async def test_count_authenticates_with_the_workload_identity_bearer( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + mock_token_exchange() + count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 12}) + ) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.error is False + assert result.total_tokens == 12 + assert count_route.calls.last.request.headers["authorization"] == "Bearer exchanged-bearer-token" + + @pytest.mark.asyncio + @respx.mock + async def test_failed_mint_becomes_a_count_error_response( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + respx.post(TOKEN_EXCHANGE_URL).mock(return_value=httpx.Response(400, json={"error": "invalid_grant"})) + count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens") + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.error is True + assert result.status_code == 401 + assert result.total_tokens == 0 + assert not count_route.called diff --git a/tests/unit/llms/vertex_ai/count_tokens/__init__.py b/tests/unit/llms/vertex_ai/count_tokens/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py new file mode 100644 index 00000000000..235c3710ea5 --- /dev/null +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -0,0 +1,507 @@ +import asyncio +import base64 +import time +import uuid +from collections.abc import Generator +from contextlib import contextmanager +from copy import deepcopy +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Thread +from typing import Final, Literal +from unittest.mock import patch + +import pytest +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm.llms.anthropic.pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + truncate_tool_name, +) +from litellm.llms.vertex_ai.count_tokens.context import ( + VertexGeminiCountContext, + VertexGeminiCountContextError, + gemini_count_contents, + inference_request_view, + vertex_gemini_count_context, +) +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + +_SCHEMA: Final[JsonValue] = { + "type": "object", + "properties": {}, + "additionalProperties": False, +} +_FUNCTION_DECLARATION: Final = { + "functionDeclarations": [ + { + "name": "lookup", + "description": "Look up a value", + "parameters": {"type": "object"}, + } + ] +} + + +@pytest.mark.parametrize( + "tool", + ( + { + "name": "lookup", + "description": "Look up a value", + "input_schema": _SCHEMA, + }, + { + "type": "function", + "function": { + "name": "lookup", + "description": "Look up a value", + "parameters": _SCHEMA, + }, + }, + { + "type": "function", + "name": "lookup", + "description": "Look up a value", + "parameters": _SCHEMA, + }, + ), + ids=("anthropic", "openai-chat", "responses"), +) +def test_function_tool_shapes_map_to_vertex_declarations(tool: JsonValue) -> None: + context: Final = vertex_gemini_count_context(system=None, tools=[tool]) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools == (_FUNCTION_DECLARATION,) + assert _SCHEMA["additionalProperties"] is False + + +def test_long_anthropic_tool_name_is_truncated_for_vertex() -> None: + name: Final = "mcp__docs__" + "a" * 70 + context: Final = vertex_gemini_count_context( + system=None, + tools=[{"name": name, "description": "Look up a value", "input_schema": _SCHEMA}], + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools == ( + { + "functionDeclarations": [ + { + "name": truncate_tool_name(name), + "description": "Look up a value", + "parameters": {"type": "object"}, + } + ] + }, + ) + assert len(truncate_tool_name(name)) == 64 + + +def test_64_character_anthropic_tool_name_is_unchanged_for_vertex() -> None: + name: Final = "mcp__docs__" + "a" * 53 + context: Final = vertex_gemini_count_context( + system=None, + tools=[{"name": name, "input_schema": _SCHEMA}], + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools == ({"functionDeclarations": [{"name": name, "parameters": {"type": "object"}}]},) + + +def test_long_openai_chat_tool_name_is_unchanged_for_vertex() -> None: + name: Final = "mcp__docs__" + "a" * 70 + context: Final = vertex_gemini_count_context( + system=None, + tools=[{"type": "function", "function": {"name": name, "parameters": _SCHEMA}}], + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools == ({"functionDeclarations": [{"name": name, "parameters": {"type": "object"}}]},) + + +def test_unknown_tools_are_omitted() -> None: + context: Final = vertex_gemini_count_context( + system=None, + tools=[{"type": "web_search_preview"}, {"type": "computer_use"}], + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools is None + + +def test_text_block_system_instruction_keeps_one_part_per_block() -> None: + context: Final = vertex_gemini_count_context( + system=[{"type": "text", "text": "First instruction."}, {"type": "text", "text": "Second instruction."}], + tools=None, + ) + + assert isinstance(context, VertexGeminiCountContext) + assert context.system_instruction == {"parts": ({"text": "First instruction."}, {"text": "Second instruction."})} + + +def test_string_system_instruction_maps_to_one_part() -> None: + context: Final = vertex_gemini_count_context(system="Follow this instruction.", tools=None) + + assert isinstance(context, VertexGeminiCountContext) + assert context.system_instruction == {"parts": ({"text": "Follow this instruction."},)} + + +def test_non_text_system_block_returns_an_error_value() -> None: + context: Final = vertex_gemini_count_context( + system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + tools=None, + ) + + assert isinstance(context, VertexGeminiCountContextError) + assert "text blocks" in context.error_message + + +@pytest.mark.parametrize(("system", "tools"), ((None, None), ("", None), ([], []))) +def test_empty_context_is_omitted(system: JsonValue | None, tools: JsonValue | None) -> None: + context: Final = vertex_gemini_count_context(system=system, tools=tools) + + assert isinstance(context, VertexGeminiCountContext) + assert context.system_instruction is None + assert context.tools is None + + +@pytest.mark.parametrize( + "tool", + ( + {"description": "Look up a value", "input_schema": _SCHEMA}, + {"type": "bash_20250124", "name": "bash"}, + ), + ids=("unnamed", "anthropic-hosted"), +) +def test_anthropic_tools_count_the_same_declarations_inference_sends(tool: dict[str, JsonValue]) -> None: + translated_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai([deepcopy(tool)]) # pyright: ignore[reportArgumentType] # raw Anthropic tool dict + inference_tools: Final = VertexGeminiConfig()._map_function( # pyright: ignore[reportPrivateUsage] # the tool mapping inference uses + value=[dict(translated) for translated in translated_tools], + optional_params={}, + ) + + context: Final = vertex_gemini_count_context(system=None, tools=[tool]) + + assert isinstance(context, VertexGeminiCountContext) + assert context.tools is not None + assert [tool_value["functionDeclarations"] for tool_value in context.tools] == [ + inference_tool["function_declarations"] for inference_tool in inference_tools + ] + + +@pytest.mark.parametrize( + "tools", + ( + "lookup", + [1], + [{"name": "lookup", "input_schema": "not-a-schema"}], + [{"type": "function", "function": {"name": "lookup", "parameters": [1]}}], + [{"name": "lookup", "input_schema": {"type": "object", "properties": {"query": "not-a-schema"}}}], + ), + ids=("not-a-list", "non-object-tool", "string-schema", "list-parameters", "malformed-property"), +) +def test_invalid_tools_return_an_error_value(tools: JsonValue) -> None: + context: Final = vertex_gemini_count_context(system=None, tools=tools) + + assert isinstance(context, VertexGeminiCountContextError) + assert "tools" in context.error_message + + +_ANTHROPIC_CONVERSATION: Final = ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize both"}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0="}}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + ], + }, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "x"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}]}, +) + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_anthropic_documents_images_and_tool_turns_become_gemini_parts( + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> None: + contents: Final = await gemini_count_contents( + _ANTHROPIC_CONVERSATION, "anthropic", "gemini-2.5-flash", custom_llm_provider, {} + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Summarize both"}, + {"inline_data": {"data": "JVBERi0=", "mime_type": "application/pdf"}}, + {"inline_data": {"data": "aW1hZ2U=", "mime_type": "image/png"}}, + ], + }, + {"role": "model", "parts": [{"function_call": {"name": "lookup", "args": {"query": "x"}}}]}, + {"role": "user", "parts": [{"function_response": {"name": "lookup", "response": {"content": "42"}}}]}, + ] + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_openai_messages_keep_their_chat_conversion(custom_llm_provider: Literal["vertex_ai", "gemini"]) -> None: + contents: Final = await gemini_count_contents( + ({"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi"}), + "openai", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + + assert contents == [ + {"role": "user", "parts": [{"text": "hello"}]}, + {"role": "model", "parts": [{"text": "hi"}]}, + ] + + +_PNG_BYTES: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +) +_DOWNLOAD_SECONDS: Final = 0.5 + + +class _SlowPngHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + time.sleep(_DOWNLOAD_SECONDS) + self.send_response(200) + self.send_header("Content-Type", "image/png") + self.send_header("Content-Length", str(len(_PNG_BYTES))) + self.end_headers() + self.wfile.write(_PNG_BYTES) + + def log_message(self, format: str, *args: object) -> None: + return None + + +@contextmanager +def _slow_png_server() -> Generator[str, None, None]: + with ThreadingHTTPServer(("127.0.0.1", 0), _SlowPngHandler) as server: + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + + +async def _event_loop_wake_delay(nap_seconds: float) -> float: + started: Final = time.monotonic() + await asyncio.sleep(nap_seconds) + return time.monotonic() - started + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_remote_image_downloads_without_blocking_the_event_loop( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with _slow_png_server() as base_url: + image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png" + wake_delay, contents = await asyncio.gather( + _event_loop_wake_delay(0.05), + gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe it"}, + {"type": "image", "source": {"type": "url", "url": image_url}}, + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ), + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Describe it"}, + {"inline_data": {"data": base64.b64encode(_PNG_BYTES).decode(), "mime_type": "image/png"}}, + ], + } + ] + assert wake_delay < _DOWNLOAD_SECONDS / 2 + + +class _MissingImageHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: + return None + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_image_that_cannot_be_fetched_returns_an_error_value( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with ThreadingHTTPServer(("127.0.0.1", 0), _MissingImageHandler) as server: + Thread(target=server.serve_forever, daemon=True).start() + try: + contents = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "url", + "url": f"http://127.0.0.1:{server.server_address[1]}/{uuid.uuid4().hex}.png", + }, + } + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + finally: + server.shutdown() + + assert isinstance(contents, VertexGeminiCountContextError) + + +@pytest.mark.parametrize( + "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") +) +def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: + view: Final = inference_request_view( + ( + {"role": "user", "content": "first question " * 50}, + {"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]}, + {"role": "user", "content": "latest question"}, + ), + system, + "anthropic", + ) + + assert view.messages == [{"role": "user", "content": "latest question"}] + assert "Summary of the earlier turns." in str(view.system) + assert "Be brief." in str(view.system) + + +def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + view: Final = inference_request_view( + ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + ), + None, + "anthropic", + ) + + assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"]) + assert [block["type"] for block in assistant_blocks] == ["text"] + assert snippet in assistant_blocks[0]["text"] + + +def test_openai_view_is_the_request_unchanged() -> None: + messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) + + view: Final = inference_request_view(messages, "Be brief.", "openai") + + assert view.messages == messages + assert view.system == "Be brief." + + +async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None: + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value="image/png", + ) as lookup: + contents: Final = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "url", "url": "gs://bucket/image-without-extension"}} + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + "vertex_ai", + {"vertex_project": "deployment-project"}, + ) + + assert lookup.call_args.kwargs.get("vertex_project") == "deployment-project" + assert not isinstance(contents, VertexGeminiCountContextError) + assert contents[0]["parts"][0] == { + "file_data": {"mime_type": "image/png", "file_uri": "gs://bucket/image-without-extension"} + } + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +def test_anthropic_replayed_web_search_result_is_counted_as_its_text( + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + messages: Final = ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + {"role": "user", "content": "Thanks."}, + ) + + view: Final = inference_request_view(messages, None, "anthropic") + contents: Final = asyncio.run( + gemini_count_contents(view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {}) + ) + + assert snippet in str(contents) diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 8407007e1d4..575e108f52f 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1,9 +1,18 @@ -from unittest.mock import patch +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Final +from unittest.mock import AsyncMock, patch +import httpx import pytest +from pydantic import JsonValue, TypeAdapter +import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler from litellm.llms.vertex_ai.common_utils import ( + VertexAITokenCounter, convert_anyof_null_to_nullable, get_vertex_location_from_url, get_vertex_project_id_from_url, @@ -14,6 +23,110 @@ from litellm.llms.vertex_ai.common_utils import ( validate_vertex_location, vertex_request_labels_from_litellm_params, ) +from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler +from litellm.types.utils import TokenCountMessagesFormat + + +@pytest.mark.asyncio +@pytest.mark.parametrize("messages_format", ("openai", "anthropic")) +async def test_vertex_gemini_rejects_non_text_system_blocks_before_request( + messages_format: TokenCountMessagesFormat, +) -> None: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-4-argon", + messages=[{"role": "user", "content": "Count this"}], + contents=None, + request_model="vertex_ai/gemini-4-argon", + system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + messages_format=messages_format, + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "text blocks" in (result.error_message or "") + + +class _VertexCountTokensPeer: + def __init__(self, status_code: int) -> None: + self.status_code: Final = status_code + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + if self.status_code == 200: + return httpx.Response(200, json={"totalTokens": 37}, request=request) + return httpx.Response( + self.status_code, + json={"error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}, + request=request, + ) + + +@contextmanager +def _vertex_count_tokens_peer(status_code: int) -> Iterator[_VertexCountTokensPeer]: + peer: Final = _VertexCountTokensPeer(status_code) + with ( + patch.object( + gemini_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ), + patch.object( + vertex_count_tokens_handler.VertexAITokenCounter, + "_ensure_access_token_async", + new=AsyncMock(return_value=("fake-token", "test-project")), + ), + ): + yield peer + + +_VERTEX_DEPLOYMENT: Final = {"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-central1"}} + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_answers_with_the_status_google_rejected_it_with() -> None: + with _vertex_count_tokens_peer(400) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "Invalid function name" in (result.error_message or "") + assert len(peer.requests) == 1 + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_rejects_an_image_url_it_refuses_to_fetch_without_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + with _vertex_count_tokens_peer(200) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}], + } + ], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "http://127.0.0.1:9/image.png" in (result.error_message or "") + assert peer.requests == () @pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"]) @@ -147,15 +260,11 @@ async def test_get_supports_system_message(): from litellm.llms.vertex_ai.common_utils import get_supports_system_message # fine-tuned vertex gemini models will specifiy they are in the /gemini spec format - result = get_supports_system_message( - model="gemini/1234567890", custom_llm_provider="vertex_ai" - ) + result = get_supports_system_message(model="gemini/1234567890", custom_llm_provider="vertex_ai") assert result == True # non-fine-tuned vertex gemini models will not specifiy they are in the /gemini spec format - result = get_supports_system_message( - model="random-model-name", custom_llm_provider="vertex_ai" - ) + result = get_supports_system_message(model="random-model-name", custom_llm_provider="vertex_ai") assert result == False @@ -228,13 +337,9 @@ def test_build_vertex_schema(): "properties": { "tags": {"items": {"type": "string"}, "type": "array"}, "metadata": {"type": "object"}, - "callbacks": { - "anyOf": [{"items": {}, "type": "array"}, {}, {"type": "null"}] - }, + "callbacks": {"anyOf": [{"items": {}, "type": "array"}, {}, {"type": "null"}]}, "run_name": {"type": "string"}, - "max_concurrency": { - "anyOf": [{"type": "integer"}, {"type": "null"}] - }, + "max_concurrency": {"anyOf": [{"type": "integer"}, {"type": "null"}]}, "recursion_limit": {"type": "integer"}, "configurable": {"type": "object"}, "run_id": { @@ -278,9 +383,7 @@ def test_build_vertex_schema(): ] }, "run_name": {"type": "string"}, - "max_concurrency": { - "anyOf": [{"type": "integer", "nullable": True}] - }, + "max_concurrency": {"anyOf": [{"type": "integer", "nullable": True}]}, "recursion_limit": {"type": "integer"}, "configurable": {"type": "object"}, "run_id": {"anyOf": [{"type": "string", "nullable": True}]}, @@ -685,9 +788,7 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix): ({}, "us-east1", "us-east1"), ], ) -def test_get_vertex_region_global_only_model( - model_cost_entry, vertex_region, expected_region -): +def test_get_vertex_region_global_only_model(model_cost_entry, vertex_region, expected_region): """Test get_vertex_region resolves region from model_cost supported_regions""" import litellm from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -699,9 +800,7 @@ def test_get_vertex_region_global_only_model( {"vertex_ai/test-model": model_cost_entry}, clear=False, ): - result = vertex_base.get_vertex_region( - vertex_region=vertex_region, model="test-model" - ) + result = vertex_base.get_vertex_region(vertex_region=vertex_region, model="test-model") assert result == expected_region @@ -819,9 +918,7 @@ def test_convert_schema_types_type_array_conversion(): assert anyof_types[1]["type"] == "number" # 4. Other properties preserved - assert ( - input_schema["properties"]["studio"]["description"] == "The studio ID or name" - ) + assert input_schema["properties"]["studio"]["description"] == "The studio ID or name" assert input_schema["required"] == ["studio"] @@ -988,7 +1085,9 @@ def test_construct_target_url_with_version_prefix(): ), ], ) -def test_construct_target_url_versionless_project_route_gets_api_version(requested_route: str, expected_url: str) -> None: +def test_construct_target_url_versionless_project_route_gets_api_version( + requested_route: str, expected_url: str +) -> None: from litellm.llms.vertex_ai.common_utils import construct_target_url target_url = construct_target_url( @@ -1121,10 +1220,7 @@ def test_fix_enum_types(): # 2. Non-string enums are removed assert "enum" not in input_schema["properties"]["maxLength"] assert "enum" not in input_schema["properties"]["enabled"] - assert ( - "enum" - not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"] - ) + assert "enum" not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"] # 3. anyOf with string type keeps enum, non-string removes it assert "enum" in input_schema["properties"]["anyOfField"]["anyOf"][0] @@ -1178,8 +1274,6 @@ def test_get_token_url(): assert "v1beta1" not in url assert "/v1/" in url - pass - @pytest.mark.asyncio async def test_vertex_ai_token_counter_routes_partner_models(): @@ -1457,6 +1551,49 @@ async def test_vertex_ai_token_counter_routes_gemini_models(): assert result.total_tokens == 50 +@pytest.mark.asyncio +async def test_vertex_gemini_count_tokens_uses_request_system_over_deployment_system_instruction() -> None: + class VertexCountTokensPeer: + def __init__(self) -> None: + self.request_body: dict[str, JsonValue] | None = None + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.request_body = TypeAdapter(dict[str, JsonValue]).validate_json(request.content) + return httpx.Response(200, json={"totalTokens": 37}, request=request) + + peer: Final = VertexCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with ( + patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), + patch.object( + vertex_count_tokens_handler.VertexAITokenCounter, + "_ensure_access_token_async", + new=AsyncMock(return_value=("fake-token", "test-project")), + ), + ): + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={ + "litellm_params": { + "vertex_project": "test-project", + "vertex_location": "us-central1", + "system_instruction": {"parts": [{"text": "deployment-level"}]}, + } + }, + request_model="vertex_ai/gemini-2.5-flash", + system="You are terse.", + ) + + assert result is not None + assert result.error is False + assert result.total_tokens == 37 + assert peer.request_body is not None + assert peer.request_body["systemInstruction"] == {"parts": [{"text": "You are terse."}]} + + @pytest.mark.asyncio async def test_vertex_ai_token_counter_converts_messages_to_contents_for_gemini(): """ @@ -1471,9 +1608,7 @@ async def test_vertex_ai_token_counter_converts_messages_to_contents_for_gemini( token_counter = VertexAITokenCounter() - with patch( - "litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens" - ) as mock_acount_tokens: + with patch("litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens") as mock_acount_tokens: mock_acount_tokens.return_value = { "totalTokens": 42, "tokenizer_used": "gemini", @@ -1515,9 +1650,7 @@ async def test_vertex_ai_token_counter_returns_none_when_api_omits_total_tokens( token_counter = VertexAITokenCounter() - with patch( - "litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens" - ) as mock_acount_tokens: + with patch("litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens") as mock_acount_tokens: mock_acount_tokens.return_value = {"tokenizer_used": "gemini"} result = await token_counter.count_tokens( @@ -1560,9 +1693,7 @@ async def test_vertex_ai_partner_model_detection(): # Test Minimax models assert VertexAIPartnerModels.is_vertex_partner_model("minimaxai/minimax-m2-maas") # Test Moonshot models - assert VertexAIPartnerModels.is_vertex_partner_model( - "moonshotai/kimi-k2-thinking-maas" - ) + assert VertexAIPartnerModels.is_vertex_partner_model("moonshotai/kimi-k2-thinking-maas") # Test Gemini models (should NOT be detected as partner model) assert not VertexAIPartnerModels.is_vertex_partner_model("gemini-1.5-pro") @@ -1593,9 +1724,7 @@ def test_vertex_ai_moonshot_uses_openai_handler(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.should_use_openai_handler( - "moonshotai/kimi-k2-thinking-maas" - ) + assert VertexAIPartnerModels.should_use_openai_handler("moonshotai/kimi-k2-thinking-maas") def test_vertex_ai_zai_uses_openai_handler(): @@ -1630,9 +1759,7 @@ def test_vertex_ai_gemma_maas_is_partner_model(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.is_vertex_partner_model( - "google/gemma-4-26b-a4b-it-maas" - ) + assert VertexAIPartnerModels.is_vertex_partner_model("google/gemma-4-26b-a4b-it-maas") def test_vertex_ai_gemma_maas_uses_openai_handler(): @@ -1643,9 +1770,7 @@ def test_vertex_ai_gemma_maas_uses_openai_handler(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.should_use_openai_handler( - "google/gemma-4-26b-a4b-it-maas" - ) + assert VertexAIPartnerModels.should_use_openai_handler("google/gemma-4-26b-a4b-it-maas") def test_vertex_ai_gemma_maas_routes_to_partner_models(): @@ -1835,12 +1960,8 @@ def test_pop_vertex_request_labels_prefers_explicit_labels_then_metadata(): def test_pop_vertex_request_labels_uses_litellm_metadata_when_metadata_absent(): optional: dict = {} - litellm_params = { - "litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}} - } - assert pop_vertex_request_labels(optional, litellm_params) == { - "team": "from_litellm_meta" - } + litellm_params = {"litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}}} + assert pop_vertex_request_labels(optional, litellm_params) == {"team": "from_litellm_meta"} def test_vertex_text_embedding_request_includes_labels_from_metadata(): @@ -1850,9 +1971,7 @@ def test_vertex_text_embedding_request_includes_labels_from_metadata(): input="hi", optional_params={}, model="text-embedding-004", - litellm_params={ - "metadata": {"requester_metadata": {"project_id": "cost-center-1"}} - }, + litellm_params={"metadata": {"requester_metadata": {"project_id": "cost-center-1"}}}, ) assert req.get("labels") == {"project_id": "cost-center-1"} diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index c0e7a4fdb47..6c7d05b6ded 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -284,7 +284,9 @@ class TestHttpExceptionDictDetail: "base_process_llm_request", new=AsyncMock(side_effect=exc), ), - patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + patch.object( + proxy_server, "proxy_logging_obj" + ) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam ): mock_logging.post_call_failure_hook = AsyncMock() response = await ep.anthropic_response( @@ -424,8 +426,6 @@ class TestErrorLogCarriesCallId: @pytest.mark.asyncio async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture): - from fastapi import HTTPException - import litellm.proxy.anthropic_endpoints.endpoints as ep import litellm.proxy.proxy_server as proxy_server from litellm.proxy._types import UserAPIKeyAuth @@ -440,13 +440,14 @@ class TestErrorLogCarriesCallId: "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]}), ), - patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block + patch.object( + proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down")) + ), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), - pytest.raises(HTTPException) as raised, ): - await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + response = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) - assert raised.value.status_code == 500 + assert response.status_code == 500 record = self._error_record(caplog) assert record.litellm_call_id == call_id assert call_id in record.getMessage() diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py new file mode 100644 index 00000000000..84a32a02a5c --- /dev/null +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -0,0 +1,190 @@ +import json +from typing import Final + +import pytest + +from litellm.proxy.common_utils.token_count_settings import ( + ProviderCountFailure, + TokenCountModel, + TokenCountSettings, + count_source_fields, + resolve_token_count_settings, +) +from litellm.types.utils import TokenCountResponse + +_OPUS_US_EAST_1: Final = TokenCountModel(model="anthropic.claude-opus-5-5", region="us-east-1") +_DEFAULT: Final = { + "token_count_framing": 97, + "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}, +} + + +def _in_us_west_2() -> str: + return "us-west-2" + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ({}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ( + {"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}}, + TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")), + ), + ( + {"token_count_model": {"model": "custom-model"}}, + TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + ), + ( + {"token_count_model": {"model": "custom-model", "region": " "}}, + TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + ), + ({"token_count_framing": 0}, TokenCountSettings(framing=0, count_model=_OPUS_US_EAST_1)), + ( + {"token_count_framing": None, "token_count_model": None}, + TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1), + ), + ( + {"token_count_framing": True, "token_count_model": "custom-model"}, + TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1), + ), + ({"token_count_framing": 78.0}, TokenCountSettings(framing=78, count_model=_OPUS_US_EAST_1)), + ({"token_count_framing": 1.5}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ({"token_count_framing": -1}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ({"token_count_framing": "78"}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ({"token_count_model": {}}, TokenCountSettings(framing=97, count_model=None)), + ({"token_count_model": ""}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ( + {"token_count_model": {"model": " ", "region": "us-east-1"}}, + TokenCountSettings(framing=97, count_model=None), + ), + ], +) +def test_config_overrides_the_cost_map_default_key_by_key( + configured: dict[str, object], expected: TokenCountSettings +) -> None: + assert ( + resolve_token_count_settings(configured=configured, default=_DEFAULT, deployment_region=_in_us_west_2) + == expected + ) + + +@pytest.mark.parametrize( + ("deployment_region", "configured", "expected_region"), + [ + ("us-west-2", {}, "us-east-1"), + ("eu-central-1", {}, None), + ("ca-central-1", {}, None), + ("us-gov-west-1", {}, None), + ( + "eu-central-1", + {"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}}, + "us-east-1", + ), + ], +) +def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geography( + deployment_region: str, configured: dict[str, object], expected_region: str | None +) -> None: + settings: Final = resolve_token_count_settings( + configured=configured, default=_DEFAULT, deployment_region=lambda: deployment_region + ) + assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region) + + +def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography() -> None: + settings: Final = resolve_token_count_settings( + configured={}, + default={"token_count_model": {"model": "global.anthropic.claude-opus-5-5", "region": "us-east-1"}}, + deployment_region=lambda: "eu-west-1", + ) + assert settings.count_model == TokenCountModel(model="global.anthropic.claude-opus-5-5", region="us-east-1") + + +def test_shipped_australian_count_region_stays_with_an_australian_deployment() -> None: + settings: Final = resolve_token_count_settings( + configured={}, + default={"token_count_model": {"model": "au.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}}, + deployment_region=lambda: "ap-southeast-2", + ) + assert settings.count_model == TokenCountModel(model="au.anthropic.claude-opus-5-5", region="ap-southeast-4") + + +def test_no_setting_anywhere_means_the_model_counts_itself_without_framing() -> None: + assert resolve_token_count_settings( + configured={}, default={}, deployment_region=_in_us_west_2 + ) == TokenCountSettings(framing=None, count_model=None) + + +@pytest.mark.parametrize( + ("value", "expected"), + [("kimi_k3", "kimi_k3"), ("cl100k_base", None), (None, None), (True, None)], +) +def test_only_a_tokenizer_litellm_ships_is_taken_from_the_cost_map(value: object, expected: str | None) -> None: + settings: Final = resolve_token_count_settings( + configured={}, default={"token_count_tokenizer": value}, deployment_region=_in_us_west_2 + ) + assert settings.tokenizer == expected + + +def _counted(source: str | None, model_used: str = "us.anthropic.claude-opus-5-5") -> TokenCountResponse: + return TokenCountResponse( + total_tokens=7, + request_model="sonnet", + model_used=model_used, + tokenizer_type="bedrock_api", + count_source=source, + ) + + +@pytest.mark.parametrize( + ("source", "expected"), + [ + (None, {}), + ("provider", {"litellm_count_source": "provider"}), + ("local_estimate", {"litellm_count_source": "local_estimate"}), + ( + "token_count_model", + {"litellm_count_source": "token_count_model", "litellm_count_model": "us.anthropic.claude-opus-5-5"}, + ), + ], +) +def test_count_source_fields_name_the_counting_model_only_when_token_count_model_counted( + source: str | None, expected: dict[str, str] +) -> None: + assert count_source_fields(_counted(source)) == expected + + +_UPSTREAM_404: Final = { + "type": "error", + "request_id": "req_upstream", + "error": {"type": "not_found_error", "message": "The model 'anthropic.claude-opus-5-5' does not exist"}, +} + + +@pytest.mark.parametrize( + ("message", "expected"), + [ + ( + json.dumps(_UPSTREAM_404), + { + **_UPSTREAM_404, + "error": { + "type": "not_found_error", + "message": "The model 'anthropic.claude-opus-5-5' does not exist. Enable eu-west-1", + }, + }, + ), + (json.dumps({"message": "Too many tokens"}), "Too many tokens. Enable eu-west-1"), + ("scripted 404", "scripted 404. Enable eu-west-1"), + ( + "Credential should be scoped to a valid region. ", + "Credential should be scoped to a valid region. Enable eu-west-1", + ), + ], +) +def test_hint_lands_in_the_upstream_error_message_and_keeps_its_envelope(message: str, expected: object) -> None: + hinted: Final = ProviderCountFailure(status_code=404, message=message).with_hint("Enable eu-west-1") + + assert (json.loads(hinted.message) if isinstance(expected, dict) else hinted.message) == expected + assert hinted.status_code == 404 diff --git a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py index 451792109de..2777d18a399 100644 --- a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py @@ -4,10 +4,11 @@ Test to verify the Google GenAI proxy API endpoints """ from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest +from litellm.types.utils import TokenCountResponse def _build_test_client(): @@ -88,9 +89,7 @@ def test_google_stream_generate_content_endpoint(): assert init_kwargs["data"]["_litellm_raw_sse_stream"] is True assert init_kwargs["data"]["_litellm_skip_openai_stream_done"] is True assert init_kwargs["data"]["model"] == "test-model" - assert init_kwargs["data"]["contents"] == [ - {"role": "user", "parts": [{"text": "Hello"}]} - ] + assert init_kwargs["data"]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] def test_google_generate_content_data_flows_through_processor(): @@ -123,9 +122,7 @@ def test_google_generate_content_data_flows_through_processor(): data = mock_init.call_args.kwargs["data"] assert data["model"] == "test-model" assert data["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] - assert data["systemInstruction"] == { - "parts": [{"text": "Your name is Doodle."}] - } + assert data["systemInstruction"] == {"parts": [{"text": "Your name is Doodle."}]} # generationConfig arrives intact here; the rename to `config` is # done downstream in route_request (see test_route_llm_request). assert data["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] @@ -160,12 +157,17 @@ def test_google_count_tokens_unchanged(): except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - fake_response = MagicMock() - fake_response.original_response = { - "totalTokens": 7, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}], - } - fake_response.total_tokens = 7 + fake_response = TokenCountResponse( + total_tokens=7, + request_model="test-model", + model_used="test-model", + tokenizer_type="gemini_api", + original_response={ + "totalTokens": 7, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}], + }, + count_source="provider", + ) with patch( "litellm.proxy.proxy_server.token_counter", @@ -180,6 +182,30 @@ def test_google_count_tokens_unchanged(): assert response.status_code == 200 body = response.json() assert body["totalTokens"] == 7 + assert body["litellm_count_source"] == "provider" + + +def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope(): + try: + client = _build_test_client() + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + with patch( + "litellm.proxy.proxy_server.token_counter", + new_callable=AsyncMock, + side_effect=RuntimeError("tokenizer unavailable"), + ): + response = client.post( + "/v1beta/models/test-model:countTokens", + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + ) + + assert response.status_code == 500 + error = response.json()["error"] + assert error["code"] == 500 + assert error["status"] == "INTERNAL" + assert "tokenizer unavailable" not in error["message"] def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs(): @@ -189,9 +215,14 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - fake_response: Final = MagicMock() - fake_response.original_response = {"input_tokens": 2167} - fake_response.total_tokens = 2167 + fake_response: Final = TokenCountResponse( + total_tokens=2167, + request_model="claude-opus-4-8", + model_used="us.anthropic.claude-opus-5-5", + tokenizer_type="bedrock_api", + original_response={"input_tokens": 2167}, + count_source="token_count_model", + ) with patch( "litellm.proxy.proxy_server.token_counter", @@ -207,3 +238,5 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs() body: Final = response.json() assert body["totalTokens"] == 2167 assert body["promptTokensDetails"] == [] + assert body["litellm_count_source"] == "token_count_model" + assert body["litellm_count_model"] == "us.anthropic.claude-opus-5-5" diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index f44b2b57ccd..54af9db84af 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -4147,6 +4147,54 @@ class TestModelInfoCostMapEchoFilter: assert "litellm_provider" not in info assert "supported_openai_params" not in info + @pytest.fixture + def bedrock_sonnet_5_5_mapped_as_on_main(self) -> Iterator[None]: + import litellm + + entry: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"}, + } + with pytest.MonkeyPatch.context() as patched: + patched.setattr(litellm, "model_cost", {**litellm.model_cost, "us.anthropic.claude-sonnet-5-5": entry}) + litellm.utils._invalidate_model_cost_lowercase_map() + yield + litellm.utils._invalidate_model_cost_lowercase_map() + + @pytest.mark.parametrize( + ("model", "field", "edited"), + [ + ("bedrock/us.anthropic.claude-sonnet-5-5", "token_count_model", {}), + ("bedrock/us.moonshotai.kimi-k3", "token_count_framing", 0), + ], + ) + def test_token_count_setting_is_stored_only_when_the_operator_changed_it( + self, model: str, field: str, edited: object, bedrock_sonnet_5_5_mapped_as_on_main: None + ) -> None: + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + entry: Final = litellm.get_model_info(model) + assert entry[field] not in (None, edited), entry + db_model: Final = Deployment( + model_name="count-settings", + litellm_params=LiteLLM_Params(model=model), + model_info=ModelInfo(id="dep-echo-count"), + ) + + echoed: Final = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count"})), + ) + changed: Final = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count", field: edited})), + ) + + assert field not in json.loads(echoed["model_info"]) + assert json.loads(changed["model_info"])[field] == edited + def test_metadata_without_a_cost_map_key_is_persisted(self): import litellm diff --git a/tests/unit/proxy/proxy_server/test_routes_utils.py b/tests/unit/proxy/proxy_server/test_routes_utils.py index 295b9e5ae8b..f93c23c8ce4 100644 --- a/tests/unit/proxy/proxy_server/test_routes_utils.py +++ b/tests/unit/proxy/proxy_server/test_routes_utils.py @@ -54,6 +54,7 @@ def test_token_counter_happy_path(client, auth_as, patched_token_counter): "request_model": "gpt-4", "model_used": "gpt-4", "tokenizer_type": "openai_tokenizer", + "count_source": "local_estimate", "original_response": None, "error": False, "error_message": None, @@ -390,8 +391,7 @@ def test_token_counter_fallback_counts_tools_system_and_anthropic_blocks(client, assert full > bare -def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, monkeypatch): - """Regression: a ``prompt`` request carrying ``tools`` but no ``messages`` still counts, because the fallback attaches tools only when counting messages (``token_counter`` rejects tools on the text path).""" +def test_token_counter_fallback_counts_a_prompt_with_tools_like_its_messages_form(client, auth_as, monkeypatch): monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False) prompt = "count the tokens in this sentence please" @@ -412,4 +412,6 @@ def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, ) assert response.status_code == 200, response.text - assert response.json()["total_tokens"] == litellm.token_counter(model="claude-fable-5", text=prompt) + assert response.json()["total_tokens"] == litellm.token_counter( + model="claude-fable-5", messages=[{"role": "user", "content": prompt}], tools=tools + ) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index e4c4ddce64f..4e0b242f80b 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2,34 +2,120 @@ # 1. Generate a Key, and use it to make a call +import asyncio +import base64 +import hashlib +import itertools +import json +import logging +import threading +import time +import traceback +from collections.abc import Callable, Coroutine, Iterator, Mapping +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import Final, NoReturn from unittest.mock import AsyncMock, MagicMock, patch +import anyio import httpx import pytest +import respx +import tiktoken from dotenv import load_dotenv -from pydantic import JsonValue +from pydantic import BaseModel, ConfigDict, JsonValue +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.processors import TemplateProcessing load_dotenv() # this file is to test litellm/proxy -from fastapi import HTTPException, Request +from fastapi import Request +from fastapi.responses import JSONResponse import litellm from litellm import Router +from litellm._logging import verbose_proxy_logger +from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler +from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + KIMI_K3_RANKS_SHA256, + KimiK3Encodings, + KimiK3TokenizerSource, + kimi_k3_template_tokens, + kimi_k3_tokens, + load_kimi_k3_encoding, +) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler +from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.anthropic_endpoints.endpoints import ( count_tokens as anthropic_count_tokens, ) -from litellm.proxy.proxy_server import token_counter -from litellm.types.utils import TokenCountResponse +from litellm.proxy.common_utils.token_count_settings import ( + ProviderCountFailure, + TokenCountModel, + TokenCountSettings, + cost_map_entry, +) +from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request +from litellm.proxy.proxy_server import ( + _estimate_or_provider_error, + _provider_token_count, + _try_provider_token_count, + count_request_tokens, + token_counter, +) +from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +from litellm.types.utils import SelectTokenizerResponse, TokenCountMessagesFormat, TokenCountResponse +def _hugging_face_offline() -> NoReturn: + raise OSError("unit tests never download the Kimi K3 tokenizer") + + +@pytest.fixture(autouse=True) +def _kimi_k3_tokenizer_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "_KIMI_K3_ENCODINGS", + KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=_hugging_face_offline, sha256=KIMI_K3_RANKS_SHA256), + retry_after_seconds=60, + ), + ) + + +@pytest.fixture +def bedrock_sonnet_5_5_mapped_as_on_main() -> Iterator[None]: + entries: Final = { + name: { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "token_count_model": {"model": count_model, "region": "us-east-1"}, + } + for name, count_model in ( + ("anthropic.claude-sonnet-5-5", "anthropic.claude-opus-5-5"), + ("us.anthropic.claude-sonnet-5-5", "us.anthropic.claude-opus-5-5"), + ) + } + with pytest.MonkeyPatch.context() as patched: + patched.setattr(litellm, "model_cost", {**litellm.model_cost, **entries}) + litellm.utils._invalidate_model_cost_lowercase_map() + yield + litellm.utils._invalidate_model_cost_lowercase_map() + @pytest.mark.asyncio async def test_vLLM_token_counting(): @@ -63,9 +149,7 @@ async def test_vLLM_token_counting(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the default tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the default tokenizer assert response.model_used == "wolfram/miquliz-120b-v2.0" @@ -98,9 +182,7 @@ async def test_token_counting_model_not_in_model_list(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the OpenAI tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the OpenAI tokenizer assert response.model_used == "special-alias" @@ -133,9 +215,7 @@ async def test_gpt_token_counting(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the OpenAI tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the OpenAI tokenizer assert response.request_model == "gpt-4" @@ -173,10 +253,8 @@ async def test_anthropic_messages_count_tokens_endpoint(): anthropic_endpoints.read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=False): - assert ( - call_endpoint == True - ), "Should be called with call_endpoint=True for Anthropic endpoint" + async def mock_token_counter(request, call_endpoint=False, messages_format="openai", output_format=None): + assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "claude-3-sonnet-20240229" assert request.messages == [{"role": "user", "content": "Hello Claude!"}] @@ -192,8 +270,8 @@ async def test_anthropic_messages_count_tokens_endpoint(): # Patch the imported token_counter function from proxy_server import litellm.proxy.proxy_server as proxy_server - original_token_counter = proxy_server.token_counter - proxy_server.token_counter = mock_token_counter + original_token_counter = proxy_server.count_request_tokens + proxy_server.count_request_tokens = mock_token_counter try: # Call the endpoint @@ -210,7 +288,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): finally: # Restore original functions anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -247,10 +325,8 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): anthropic_endpoints.read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=True): - assert ( - call_endpoint == True - ), "Should be called with call_endpoint=True for Anthropic endpoint" + async def mock_token_counter(request, call_endpoint=True, messages_format="openai", output_format=None): + assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "gpt-4" assert request.messages == [{"role": "user", "content": "Hello GPT!"}] @@ -266,8 +342,8 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): # Patch the imported token_counter function from proxy_server import litellm.proxy.proxy_server as proxy_server - original_token_counter = proxy_server.token_counter - proxy_server.token_counter = mock_token_counter + original_token_counter = proxy_server.count_request_tokens + proxy_server.count_request_tokens = mock_token_counter try: # Call the endpoint @@ -284,77 +360,79 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): finally: # Restore original functions anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio -async def test_internal_token_counter_anthropic_provider_detection(): - """ - Test that the internal token_counter correctly detects Anthropic providers - and handles the from_anthropic_endpoint flag appropriately - """ - - # Test with Anthropic provider - llm_router = Router( - model_list=[ - { - "model_name": "claude-test", - "litellm_params": { - "model": "anthropic/claude-3-sonnet-20240229", - "api_key": "test-key", - }, - } - ] +@pytest.mark.parametrize("status_code", (400, 401, 403, 404, 429, 500)) +@pytest.mark.parametrize( + ("model", "count_url"), + ( + ("anthropic/claude-sonnet-5-5", "https://api.anthropic.com/v1/messages/count_tokens"), + ("openai/gpt-6.1-sol", "https://api.openai.com/v1/responses/input_tokens"), + ), +) +async def test_provider_errors_return_labelled_local_estimates( + monkeypatch: pytest.MonkeyPatch, status_code: int, model: str, count_url: str +) -> None: + messages: Final = [ + {"role": "user", "content": "Count the whole request locally when the count API is unavailable."} + ] + router: Final = Router( + model_list=[{"model_name": "counted", "litellm_params": {"model": model, "api_key": "test-count-key"}}] ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "disable_token_counter", False) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") - setattr(litellm.proxy.proxy_server, "llm_router", llm_router) + with respx.mock(assert_all_called=True) as upstream: + upstream.post(count_url).mock( + return_value=httpx.Response( + status_code, json={"type": "error", "error": {"type": "api_error", "message": "Count unavailable"}} + ) + ) + counted: Final = await token_counter( + request=TokenCountRequest(model="counted", messages=messages), call_endpoint=True + ) - # Test with is_direct_request=False (simulating call from Anthropic endpoint) - response = await token_counter( - request=TokenCountRequest( - model="claude-test", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, + assert counted.total_tokens == litellm.token_counter(model=model.split("/", 1)[1], messages=messages) + assert counted.count_source == "local_estimate" + assert counted.error is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "part", + ( + {"fileData": {"mimeType": "image/png", "fileUri": "gs://bucket/image.png"}}, + {"inlineData": {"mimeType": "image/png", "data": "aW1hZ2U="}}, + {"text": "Reasoning about the request", "thought": True}, + ), +) +async def test_google_provider_errors_preserve_legacy_translated_message_estimates( + monkeypatch: pytest.MonkeyPatch, part: dict[str, JsonValue] +) -> None: + model: Final = "openai/gpt-6.1-sol" + request: Final = _google_count_tokens_request( + "counted", {"contents": [{"role": "user", "parts": [{"text": "Describe this"}, part]}]} ) - - print("Anthropic provider test response:", response) - - # Verify response structure - assert response.request_model == "claude-test" - assert response.model_used == "claude-3-sonnet-20240229" - assert response.total_tokens > 0 - - # Test with non-Anthropic provider - llm_router = Router( - model_list=[ - { - "model_name": "gpt-test", - "litellm_params": { - "model": "gpt-4", - }, - } - ] + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "counted", "litellm_params": {"model": model, "api_key": "test-key"}}]), ) - - setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - - # Test with is_direct_request=False but non-Anthropic provider - response = await token_counter( - request=TokenCountRequest( - model="gpt-test", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) - - print("Non-Anthropic provider test response:", response) - - # Verify response structure - assert response.request_model == "gpt-test" - assert response.model_used == "gpt-4" - assert response.total_tokens > 0 - assert response.tokenizer_type == "openai_tokenizer" # Should use LiteLLM tokenizer + monkeypatch.setattr(litellm, "disable_token_counter", False) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + with respx.mock(assert_all_called=True) as upstream: + upstream.post("https://api.openai.com/v1/responses/input_tokens").mock( + return_value=httpx.Response(400, json={"error": {"message": "Count unavailable"}}) + ) + counted: Final = await token_counter(request=request, call_endpoint=True) + assert counted.total_tokens == litellm.token_counter(model="gpt-6.1-sol", messages=request.messages) + assert counted.count_source == "local_estimate" + assert counted.error is False @pytest.mark.asyncio @@ -364,12 +442,14 @@ async def test_anthropic_endpoint_error_handling(): """ from unittest.mock import MagicMock - from fastapi import HTTPException, Request + from fastapi import Request from litellm.proxy.anthropic_endpoints.endpoints import count_tokens # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} + mock_request.scope = {} mock_user_api_key_dict = MagicMock() # Test missing model parameter @@ -387,12 +467,13 @@ async def test_anthropic_endpoint_error_handling(): anthropic_endpoints.read_request_body = mock_read_request_body try: - # Should raise HTTPException for missing model - with pytest.raises(HTTPException) as exc_info: - await count_tokens(mock_request, mock_user_api_key_dict) + response = await count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 400 - assert "model parameter is required" in str(exc_info.value.detail) + assert response.status_code == 400 + assert json.loads(response.body) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "model parameter is required"}, + } print("✅ Error handling test passed!") @@ -411,9 +492,7 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -470,9 +549,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -487,7 +564,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): mock_router.model_list = [ { "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, + "litellm_params": {"model": "databricks/gpt-4"}, "model_info": {}, } ] @@ -496,7 +573,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): mock_router.async_get_available_deployment = AsyncMock( return_value={ "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, + "litellm_params": {"model": "databricks/gpt-4"}, "model_info": {}, } ) @@ -531,9 +608,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -548,9 +623,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): mock_router.model_list = [ { "model_name": "claude-3-5-sonnet", - "litellm_params": { - "model": "anthropic/claude-3-5-sonnet-20241022" - }, + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}, "model_info": {}, } ] @@ -559,9 +632,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): mock_router.async_get_available_deployment = AsyncMock( return_value={ "model_name": "claude-3-5-sonnet", - "litellm_params": { - "model": "anthropic/claude-3-5-sonnet-20241022" - }, + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}, "model_info": {}, } ) @@ -596,9 +667,7 @@ async def test_factory_registration(): assert counter is not None # Create test deployments - anthropic_deployment = { - "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"} - } + anthropic_deployment = {"litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}} non_anthropic_deployment = {"litellm_params": {"model": "openai/gpt-4"}} @@ -634,9 +703,7 @@ async def test_bedrock_count_tokens_endpoint(): model_list=[ { "model_name": "claude-bedrock", - "litellm_params": { - "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0" - }, + "litellm_params": {"model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0"}, } ] ) @@ -650,9 +717,7 @@ async def test_bedrock_count_tokens_endpoint(): } # Test the mock handler directly to verify correct parameter extraction - await mock_count_tokens_handler( - request_data, {}, "anthropic.claude-3-sonnet-20240229-v1:0" - ) + await mock_count_tokens_handler(request_data, {}, "anthropic.claude-3-sonnet-20240229-v1:0") @pytest.mark.asyncio @@ -714,10 +779,7 @@ async def test_vertex_ai_anthropic_token_counting(): assert call_args is not None assert call_args.kwargs["model"] == "claude-3-5-sonnet-20241022" assert "messages" in call_args.kwargs["request_data"] - assert ( - call_args.kwargs["request_data"]["messages"][0]["content"] - == "Hello Claude on Vertex AI! How are you?" - ) + assert call_args.kwargs["request_data"]["messages"][0]["content"] == "Hello Claude on Vertex AI! How are you?" # Validate response structure assert response.model_used == "claude-3-5-sonnet-20241022" @@ -750,9 +812,7 @@ def test_vertex_ai_partner_models_token_counting_endpoint(vertex_location): if vertex_location == "global": assert endpoint.startswith("https://aiplatform.googleapis.com") else: - assert endpoint.startswith( - f"https://{vertex_location}-aiplatform.googleapis.com" - ) + assert endpoint.startswith(f"https://{vertex_location}-aiplatform.googleapis.com") class _FailingRuntimeHandler(BedrockCountTokensHandler): @@ -838,20 +898,14 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): with patch.object(handler, "validate_count_tokens_request"): with patch.object(handler, "_get_aws_region_name", return_value="us-west-2"): - with patch.object( - handler, "transform_anthropic_to_bedrock_count_tokens", return_value={} - ): + with patch.object(handler, "transform_anthropic_to_bedrock_count_tokens", return_value={}): with patch.object( handler, "get_bedrock_count_tokens_endpoint", return_value="https://example.com", ): - with patch.object( - handler, "_sign_request", return_value=({}, "{}") - ): - with patch( - "litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client" - ) as mock_client: + with patch.object(handler, "_sign_request", return_value=({}, "{}")): + with patch("litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client") as mock_client: mock_async_client = AsyncMock() mock_async_client.post = AsyncMock(side_effect=http_error) mock_client.return_value = mock_async_client @@ -860,9 +914,7 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): await handler.handle_count_tokens_request( request_data={ "model": "test", - "messages": [ - {"role": "user", "content": "hello"} - ], + "messages": [{"role": "user", "content": "hello"}], }, litellm_params={}, resolved_model="anthropic.claude-3-sonnet", @@ -870,81 +922,9 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): assert exc_info.value.status_code == 403 # Message should be the raw response text - assert ( - exc_info.value.message - == "Forbidden - Invalid credentials" - ) + assert exc_info.value.message == "Forbidden - Invalid credentials" -@pytest.mark.asyncio -async def test_token_counter_httpx_status_error_raises_proxy_exception(): - """ - When provider_counter.count_tokens() raises httpx.HTTPStatusError, - the token_counter endpoint should catch it and raise a ProxyException - with the upstream status code and error message. - """ - - upstream_status = 429 - upstream_message = "Rate limit exceeded" - response = httpx.Response( - status_code=upstream_status, - request=httpx.Request("POST", "https://provider.example.com/count"), - ) - http_error = httpx.HTTPStatusError( - message=upstream_message, - request=response.request, - response=response, - ) - - mock_counter = MagicMock() - mock_counter.should_use_token_counting_api.return_value = True - mock_counter.count_tokens = AsyncMock(side_effect=http_error) - - # Save originals - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) - original_router = litellm.proxy.proxy_server.llm_router - - try: - - def mock_get_provider_token_counter(deployment, model_to_use): - return (mock_counter, "claude-4-6-sonnet", "vertex_ai") - - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) - - mock_router = MagicMock() - mock_router.async_get_available_deployment = AsyncMock( - return_value={ - "litellm_params": { - "model": "vertex_ai/claude-4-6-sonnet", - "api_key": "fake-key", - }, - "model_info": {}, - } - ) - litellm.proxy.proxy_server.llm_router = mock_router - - with pytest.raises(ProxyException) as exc_info: - await token_counter( - request=TokenCountRequest( - model="claude-4-6-sonnet", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) - - assert exc_info.value.code == str(upstream_status) - assert upstream_message in exc_info.value.message - assert exc_info.value.type == "token_counting_error" - assert exc_info.value.param == "model" - finally: - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) - litellm.proxy.proxy_server.llm_router = original_router @pytest.mark.asyncio @@ -979,9 +959,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): # Save original value and function original_disable = litellm.disable_token_counter - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter try: litellm.disable_token_counter = True @@ -995,9 +973,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): def mock_get_provider_token_counter(deployment, model_to_use): return (mock_counter, "anthropic.claude-3-sonnet", "bedrock") - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter with pytest.raises(ProxyException) as exc_info: await token_counter( @@ -1012,81 +988,9 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): assert "Rate limit exceeded" in exc_info.value.message finally: litellm.disable_token_counter = original_disable - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter -@pytest.mark.asyncio -async def test_proxy_token_counter_error_falls_back_when_enabled(): - """ - Test that proxy token_counter falls back to local tokenizer when disable_token_counter=False - and provider returns an error response. - """ - # Create error response - error_response = TokenCountResponse( - total_tokens=0, - request_model="bedrock/anthropic.claude-3-sonnet", - model_used="anthropic.claude-3-sonnet", - tokenizer_type="bedrock_api", - error=True, - error_message="Rate limit exceeded", - status_code=429, - ) - - # Create mock router that returns a deployment - mock_deployment = { - "litellm_params": { - "model": "bedrock/anthropic.claude-3-sonnet", - }, - "model_info": {}, - } - - mock_router = MagicMock() - mock_router.async_get_available_deployment = AsyncMock(return_value=mock_deployment) - - setattr(litellm.proxy.proxy_server, "llm_router", mock_router) - - # Save original value and function - original_disable = litellm.disable_token_counter - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) - - try: - litellm.disable_token_counter = False - - # Create a mock counter that returns an error response - mock_counter = MagicMock(spec=BedrockTokenCounter) - mock_counter.should_use_token_counting_api.return_value = True - mock_counter.count_tokens = AsyncMock(return_value=error_response) - - # Replace the function directly - def mock_get_provider_token_counter(deployment, model_to_use): - return (mock_counter, "anthropic.claude-3-sonnet", "bedrock") - - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) - - # Should not raise, should fall back to local tokenizer - result = await token_counter( - request=TokenCountRequest( - model="claude-bedrock", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) - - # Should have used the fallback tokenizer - assert result.error is False - assert result.total_tokens > 0 - assert result.tokenizer_type != "bedrock_api" - finally: - litellm.disable_token_counter = original_disable - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) @pytest.mark.asyncio @@ -1099,6 +1003,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1112,10 +1017,10 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): original_read_request_body = anthropic_endpoints.read_request_body anthropic_endpoints.read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with Bedrock-style error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message='{"detail":{"message":"Input is too long for requested model."}}', type="token_counting_error", @@ -1123,23 +1028,20 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): code=400, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - # Verify HTTP status code is correct - assert exc_info.value.status_code == 400 - - # Verify error is in Anthropic format - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "invalid_request_error" assert detail["error"]["message"] == "Input is too long for requested model." finally: anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -1151,6 +1053,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): import litellm.proxy.proxy_server as proxy_server mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1164,10 +1067,10 @@ async def test_anthropic_endpoint_403_permission_error_format(): original_read_request_body = anthropic_endpoints.read_request_body anthropic_endpoints.read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 403 error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message='{"Message":"Bearer Token has expired"}', type="token_counting_error", @@ -1175,21 +1078,20 @@ async def test_anthropic_endpoint_403_permission_error_format(): code=403, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 403 - - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 403 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "permission_error" assert detail["error"]["message"] == "Bearer Token has expired" finally: anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -1201,6 +1103,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): import litellm.proxy.proxy_server as proxy_server mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1214,10 +1117,10 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): original_read_request_body = anthropic_endpoints.read_request_body anthropic_endpoints.read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 429 error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message="Rate limit exceeded", type="token_counting_error", @@ -1225,18 +1128,1537 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): code=429, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 429 - - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 429 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "rate_limit_error" assert detail["error"]["message"] == "Rate limit exceeded" finally: anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter + + +class _ScriptedCounter(BaseTokenCounter): + def __init__(self, replies: dict[str, TokenCountResponse], burst: int = 1) -> None: + self._replies = replies + self._burst = burst + self._burst_arrived = asyncio.Event() + self.calls: list[tuple[str, object, object, object]] = [] + self.regions: list[object] = [] + + def should_use_token_counting_api(self, custom_llm_provider: str | None = None) -> bool: + return True + + async def count_tokens( + self, + model_to_use: str, + messages: object, + contents: object, + deployment: object = None, + request_model: str = "", + tools: object = None, + system: object = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, + context_management: object = None, + ) -> TokenCountResponse | None: + self.calls.append((model_to_use, messages, tools, system)) + self.regions.append(None if deployment is None else deployment["litellm_params"].get("aws_region_name")) + if len(self.calls) >= self._burst: + self._burst_arrived.set() + await self._burst_arrived.wait() + return self._replies[model_to_use] + + +def _provider_reply( + model: str, total_tokens: int = 0, status_code: int | None = None, rejected_before_provider_call: bool = False +) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=total_tokens, + request_model="sonnet-5-5", + model_used=model, + tokenizer_type="bedrock_api", + error=status_code is not None, + error_message=None if status_code is None else f"scripted {status_code} for {model}", + status_code=status_code, + rejected_before_provider_call=rejected_before_provider_call, + ) + + +_COUNT_REQUEST = TokenCountRequest( + model="sonnet-5-5", + messages=[{"role": "user", "content": "hello"}], + tools=[{"name": "get_quote", "input_schema": {"type": "object"}}], + system="be brief", +) + + +_SONNET_DEPLOYMENT: Final = MappingProxyType( + {"litellm_params": {"model": "bedrock/us.anthropic.claude-sonnet-5-5", "aws_region_name": "us-west-2"}} +) +_OPUS_IN_US_EAST_1: Final = TokenCountModel(model="us.anthropic.claude-opus-5-5", region="us-east-1") + + +async def _ladder( + counter: _ScriptedCounter, settings: TokenCountSettings, local_estimate_allowed: bool = True +) -> TokenCountResponse | ProviderCountFailure | None: + return await _provider_token_count( + count_with=partial( + _try_provider_token_count, + counter, + "bedrock", + messages=_COUNT_REQUEST.messages, + contents=_COUNT_REQUEST.contents, + deployment=_SONNET_DEPLOYMENT, + request_model=_COUNT_REQUEST.model, + tools=_COUNT_REQUEST.tools, + system=_COUNT_REQUEST.system, + ), + model_to_use="us.anthropic.claude-sonnet-5-5", + settings=settings, + local_estimate_allowed=local_estimate_allowed, + ) + + +@pytest.mark.asyncio +async def test_model_without_token_count_model_counts_itself_in_the_deployment_region(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=447)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=10, count_model=None)) + + assert result is not None + assert (result.total_tokens, result.count_source) == (447, "provider") + assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-sonnet-5-5"], ["us-west-2"]) + + +@pytest.mark.parametrize( + ("count_model", "region"), + [ + (_OPUS_IN_US_EAST_1, "us-east-1"), + (TokenCountModel(model="us.anthropic.claude-opus-5-5", region=None), "us-west-2"), + ], +) +@pytest.mark.asyncio +async def test_token_count_model_counts_the_same_request_in_its_own_region_without_trying_the_model( + count_model: TokenCountModel, region: str +): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=None, count_model=count_model)) + + assert result is not None + assert (result.total_tokens, result.count_source, result.model_used) == ( + 447, + "token_count_model", + "us.anthropic.claude-opus-5-5", + ) + assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-opus-5-5"], [region]) + assert counter.calls[0][1:] == (_COUNT_REQUEST.messages, _COUNT_REQUEST.tools, _COUNT_REQUEST.system) + assert _SONNET_DEPLOYMENT["litellm_params"]["aws_region_name"] == "us-west-2" + + +@pytest.mark.parametrize("status_code", (400, 401, 403, 404, 429, 500)) +@pytest.mark.parametrize("framing", (None, 0, 10)) +@pytest.mark.asyncio +async def test_provider_count_errors_preserve_legacy_local_fallback(status_code: int, framing: int | None): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=status_code)} + ) + + result: Final = await _ladder(counter, TokenCountSettings(framing=framing, count_model=None)) + + assert isinstance(result, ProviderCountFailure) + assert (result.status_code, result.message) == (status_code, f"scripted {status_code} for sonnet") + + +@pytest.mark.parametrize("status_code", [400, 403, 404, 500]) +@pytest.mark.asyncio +async def test_token_count_model_failure_names_its_region_so_the_user_can_enable_it(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + result: Final = await _ladder(counter, TokenCountSettings(framing=None, count_model=_OPUS_IN_US_EAST_1)) + + assert isinstance(result, ProviderCountFailure) + assert result.status_code == status_code + assert result.message.startswith(f"scripted {status_code} for opus") + assert "us.anthropic.claude-opus-5-5 in us-east-1" in result.message + assert len(counter.calls) == 1 + + +@pytest.mark.asyncio +async def test_token_count_model_without_a_region_preserves_its_error_for_a_failed_local_estimate(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=404)}) + + result: Final = await _ladder( + counter, + TokenCountSettings(framing=None, count_model=TokenCountModel(model="us.anthropic.claude-opus-5-5", region=None)), + ) + + assert isinstance(result, ProviderCountFailure) + assert (result.status_code, result.message) == (404, "scripted 404 for opus") + + +@pytest.mark.asyncio +async def test_uncountable_model_without_framing_preserves_legacy_local_fallback(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=404)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=None, count_model=None)) + + assert isinstance(result, ProviderCountFailure) + assert (result.status_code, result.message) == (404, "scripted 404 for sonnet") + + +@pytest.mark.asyncio +async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=None)) + + assert isinstance(result, ProviderCountFailure) + assert len(counter.calls) == 1 + + +@pytest.mark.parametrize("status_code", [400, 404]) +@pytest.mark.asyncio +async def test_token_count_model_failure_allows_a_framed_local_estimate(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + result: Final = await _ladder(counter, TokenCountSettings(framing=78, count_model=_OPUS_IN_US_EAST_1)) + + assert isinstance(result, ProviderCountFailure) + assert result.status_code == status_code + assert "us.anthropic.claude-opus-5-5 in us-east-1" in result.message + + +@pytest.mark.asyncio +async def test_uncountable_model_with_framing_raises_when_local_counting_is_disabled(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=10, count_model=None), local_estimate_allowed=False) + + assert raised.value.code == "400" + + +@pytest.mark.asyncio +async def test_rejected_provider_conversion_still_allows_a_local_estimate(): + counter: Final = _ScriptedCounter( + { + "us.anthropic.claude-sonnet-5-5": _provider_reply( + "sonnet", status_code=400, rejected_before_provider_call=True + ) + } + ) + + result: Final = await _ladder(counter, TokenCountSettings(framing=100, count_model=None)) + + assert isinstance(result, ProviderCountFailure) + assert (result.status_code, result.message) == (400, "scripted 400 for sonnet") + + +@pytest.mark.parametrize("model", ["openai/gpt-5.4-mini", "gemini/gemini-3.8-flash"]) +@pytest.mark.asyncio +async def test_framing_keeps_the_400_for_blocks_the_provider_counter_cannot_convert( + monkeypatch: pytest.MonkeyPatch, model: str +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "framed", + "litellm_params": {"model": model, "api_key": "sk-unused-framed"}, + "model_info": {"token_count_framing": 100}, + } + ] + ), + ) + search_result: Final = {"type": "search_result", "source": "https://example.com", "title": "Notes", "content": []} + + with pytest.raises(ProxyException) as raised: + await count_request_tokens( + request=TokenCountRequest( + model="framed", + messages=[{"role": "user", "content": [{"type": "text", "text": "Summarize."}, search_result]}], + ), + call_endpoint=True, + messages_format="anthropic", + ) + + assert raised.value.code == "400" + assert "search_result" in raised.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "unreadable_body", + [ + pytest.param( + { + "messages": [ + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "new"}, + }, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_01", "content": []}, + ], + }, + ] + }, + id="web_search_blocks", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "What is new?"}], + "tools": [{"name": "search", "input_schema": {"type": "object", "properties": {"query": True}}}], + }, + id="boolean_subschema_tool", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "What is new?"}], + "tools": [ + {"name": "search", "input_schema": {"type": "object", "properties": {"tags": {"type": "array"}}}} + ], + }, + id="array_schema_without_items", + ), + ], +) +async def test_framed_estimate_that_cannot_read_the_body_answers_the_provider_error( + monkeypatch: pytest.MonkeyPatch, unreadable_body: dict[str, list[dict[str, object]]] +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + with pytest.raises(ProxyException) as raised: + await token_counter(request=TokenCountRequest(model="kimi", **unreadable_body), call_endpoint=True) + readable: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for kimi") + assert readable.count_source == "local_estimate" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "contents", + [ + pytest.param([{"role": "user", "parts": [{"text": "What is new in the release notes?"}]}], id="text"), + pytest.param( + [ + {"role": "user", "parts": [{"text": "Weather in Paris?"}]}, + { + "role": "model", + "parts": [ + {"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}, "thoughtSignature": "c2ln"} + ], + }, + { + "role": "user", + "parts": [{"functionResponse": {"name": "get_weather", "response": {"forecast": "Sun"}}}], + }, + ], + id="sdk_agent_loop", + ), + pytest.param( + [ + {"role": "user", "parts": [{"text": "Hi"}]}, + {"role": "model", "parts": [{"text": "Hello!", "thoughtSignature": "c2ln"}]}, + {"role": "user", "parts": [{"text": "What is new?"}]}, + ], + id="sdk_chat_with_a_thought_signature", + ), + pytest.param([], id="empty"), + ], +) +async def test_framed_estimate_counts_google_contents_like_the_same_conversation_sent_as_messages( + monkeypatch: pytest.MonkeyPatch, contents: list[dict[str, object]] +): + _framed_kimi(monkeypatch) + google_request: Final = _google_count_tokens_request("kimi", {"contents": contents}) + + from_google: Final = await token_counter(request=google_request, call_endpoint=True) + from_messages: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=google_request.messages), call_endpoint=True + ) + + assert from_google.count_source == "local_estimate" + assert from_google.total_tokens == from_messages.total_tokens + + +@pytest.mark.asyncio +async def test_framed_deployment_whose_custom_tokenizer_cannot_load_answers_the_provider_error( + monkeypatch: pytest.MonkeyPatch, +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "kimi", + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": {"custom_tokenizer": {"identifier": "moonshotai/kimi-k3", "revision": "main"}}, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for kimi") + + +def _framed_kimis_with_custom_tokenizers( + monkeypatch: pytest.MonkeyPatch, tokenizers: Mapping[str, tuple[str, str]], burst: int = 1 +) -> None: + counter: Final = _ScriptedCounter( + { + "us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400), + "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=12), + }, + burst=burst, + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": {"model": "bedrock/us.anthropic.claude-sonnet-5-5"}, + "model_info": {"token_count_model": {}}, + }, + *( + { + "model_name": name, + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": { + "custom_tokenizer": {"identifier": identifier, "revision": revision, "auth_token": None} + }, + } + for name, (identifier, revision) in tokenizers.items() + ), + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + +def _framed_kimi_with_custom_tokenizer(monkeypatch: pytest.MonkeyPatch, identifier: str) -> None: + _framed_kimis_with_custom_tokenizers(monkeypatch, {"kimi": (identifier, "main")}) + + +async def _count_status(model: str = "kimi") -> str: + try: + response: Final = await token_counter( + request=TokenCountRequest(model=model, messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + except ProxyException as e: + return e.code + return response.count_source or "" + + +class _SkippableClockLoop(asyncio.SelectorEventLoop): + skipped_seconds: float = 0.0 + + def time(self) -> float: + return super().time() + self.skipped_seconds + + def skip(self, seconds: float) -> None: + self.skipped_seconds += seconds + + +def _run_on_skippable_clock( + scenario: Callable[[_SkippableClockLoop], Coroutine[object, object, tuple[str, ...]]], +) -> tuple[str, ...]: + loop: Final = _SkippableClockLoop() + try: + return loop.run_until_complete(scenario(loop)) + finally: + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + +async def _until_finished(tasks: tuple[asyncio.Future[str], ...], count: int) -> None: + for finished in itertools.islice(asyncio.as_completed(tasks), count): + await finished + + +def _tokenizer_that_panics_while_encoding() -> SelectTokenizerResponse: + template: Final = Tokenizer(WordLevel({"[UNK]": 0, "": 1}, unk_token="[UNK]")) + template.pre_tokenizer = Whitespace() + template.post_processor = TemplateProcessing(single=" $A", special_tokens=[("", 1)]) + spec: Final = json.loads(template.to_str()) + return litellm.utils.create_tokenizer( + json.dumps({**spec, "post_processor": {**spec["post_processor"], "special_tokens": {}}}) + ) + + +def test_hanging_tokenizer_download_answers_the_provider_error_after_the_wait_with_one_download_while_it_runs( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/hanging-tokenizer") + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + downloads: Final[list[str]] = [] + + def hanging_download(model: str, custom_tokenizer: object = None) -> NoReturn: + downloads.append(model) + download_started.set() + hub_released.wait(10) + raise OSError("hub unreachable") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hanging_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e-9) + + async def burst_past_the_wait_then_one_more(loop: _SkippableClockLoop) -> tuple[str, ...]: + try: + burst: Final = asyncio.gather(*(_count_status() for _ in range(60))) + await asyncio.to_thread(download_started.wait, 10) + loop.skip(31) + answered: Final = await asyncio.wait_for(burst, 8) + later: Final = await asyncio.wait_for(_count_status(), 8) + finally: + hub_released.set() + return (*answered, later) + + assert set(_run_on_skippable_clock(burst_past_the_wait_then_one_more)) == {"400"} + assert len(downloads) == 1 + + +@pytest.mark.asyncio +async def test_framed_count_waits_for_a_tokenizer_download_that_finishes_within_the_wait( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/slow-tokenizer") + real_select_tokenizer: Final = litellm.utils._select_tokenizer + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + + def slow_download(model: str, custom_tokenizer: object = None): + download_started.set() + hub_released.wait(10) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", slow_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + count: Final = asyncio.create_task(_count_status()) + try: + await asyncio.to_thread(download_started.wait, 10) + finally: + hub_released.set() + + assert await asyncio.wait_for(count, 10) == "local_estimate" + + +@pytest.mark.asyncio +async def test_failed_tokenizer_download_is_reused_until_the_retry_window_passes(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/flaky-tokenizer") + real_select_tokenizer: Final = litellm.utils._select_tokenizer + downloads: Final[list[str]] = [] + + def hub_down_once(model: str, custom_tokenizer: object = None): + downloads.append(model) + if len(downloads) == 1: + raise OSError("hub unreachable") + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hub_down_once) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e9) + within_window: Final = [await _count_status() for _ in range(3)] + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e-9) + after_window: Final = await _count_status() + + assert within_window == ["400", "400", "400"] + assert after_window == "local_estimate" + assert len(downloads) == 2 + + +def test_failed_tokenizer_load_is_reused_for_the_window_after_it_fails(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/slow-failing-tokenizer") + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + downloads: Final[list[str]] = [] + + def slow_failing_download(model: str, custom_tokenizer: object = None) -> NoReturn: + downloads.append(model) + download_started.set() + hub_released.wait(10) + raise OSError("hub unreachable") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", slow_failing_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 1000) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 60) + + async def fail_after_longer_than_the_window_then_retry(loop: _SkippableClockLoop) -> tuple[str, ...]: + first: Final = asyncio.ensure_future(_count_status()) + try: + await asyncio.to_thread(download_started.wait, 10) + loop.skip(120) + finally: + hub_released.set() + return await asyncio.wait_for(first, 8), await asyncio.wait_for(_count_status(), 8) + + assert _run_on_skippable_clock(fail_after_longer_than_the_window_then_retry) == ("400", "400") + assert len(downloads) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "second_tokenizer", + [("lit-tests/tokenizer-b", "main"), ("lit-tests/tokenizer-a", "v2")], + ids=["identifier", "revision"], +) +async def test_framed_deployments_with_different_custom_tokenizers_each_load_their_own( + monkeypatch: pytest.MonkeyPatch, second_tokenizer: tuple[str, str] +): + _framed_kimis_with_custom_tokenizers( + monkeypatch, {"kimi-a": ("lit-tests/tokenizer-a", "main"), "kimi-b": second_tokenizer} + ) + real_select_tokenizer: Final = litellm.utils._select_tokenizer + loaded: Final[list[object]] = [] + + def load(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + loaded.append(custom_tokenizer) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", load) + + statuses: Final = [await _count_status(name) for name in ("kimi-a", "kimi-b")] + + assert statuses == ["local_estimate", "local_estimate"] + assert loaded == [ + {"identifier": "lit-tests/tokenizer-a", "revision": "main", "auth_token": None}, + {"identifier": second_tokenizer[0], "revision": second_tokenizer[1], "auth_token": None}, + ] + + +@pytest.mark.asyncio +async def test_tokenizer_loads_past_the_cap_answer_the_provider_error_without_stalling_other_counts( + monkeypatch: pytest.MonkeyPatch, +): + tokenizers: Final = {f"kimi-{n}": (f"lit-tests/unreachable-tokenizer-{n}", "main") for n in range(5)} + _framed_kimis_with_custom_tokenizers(monkeypatch, tokenizers, burst=len(tokenizers)) + real_select_tokenizer: Final = litellm.utils._select_tokenizer + hub_released: Final = threading.Event() + downloads: Final[list[object]] = [] + + def held_download(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + downloads.append(custom_tokenizer) + hub_released.wait(10) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", held_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 2) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + monkeypatch.setattr(anyio.to_thread.current_default_thread_limiter(), "total_tokens", 2) + framed: Final = tuple(asyncio.ensure_future(_count_status(name)) for name in tokenizers) + try: + await asyncio.wait_for(_until_finished(framed, 3), 8) + unrelated: Final = await asyncio.wait_for(_count_status("sonnet"), 5) + finally: + hub_released.set() + statuses: Final = await asyncio.wait_for(asyncio.gather(*framed), 8) + + assert sorted(statuses) == ["400"] * 3 + ["local_estimate"] * 2 + assert unrelated == "provider" + assert len(downloads) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("panics_while", ["loading", "counting"]) +async def test_framed_count_answers_the_provider_error_when_the_tokenizer_panics( + monkeypatch: pytest.MonkeyPatch, panics_while: str +): + _framed_kimi_with_custom_tokenizer(monkeypatch, f"lit-tests/panicking-tokenizer-{panics_while}") + panicking: Final = _tokenizer_that_panics_while_encoding() + + def load(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + if panics_while == "loading": + panicking["tokenizer"].encode("hello world") + return panicking + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", load) + + assert await _count_status() == "400" + + +def _framed_kimi(monkeypatch: pytest.MonkeyPatch) -> None: + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + +@pytest.mark.asyncio +async def test_local_count_loads_the_tokenizer_again_after_a_failed_load(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/local-route-tokenizer") + real_select_tokenizer: Final = litellm.utils._select_tokenizer + downloads: Final[list[str]] = [] + + def hub_down_once(model: str, custom_tokenizer: object = None): + downloads.append(model) + if len(downloads) == 1: + raise OSError("hub unreachable") + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hub_down_once) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e9) + request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]) + + with pytest.raises(OSError, match="hub unreachable"): + await token_counter(request=request) + retried: Final = await token_counter(request=request) + + assert retried.count_source == "local_estimate" and retried.total_tokens > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("reuse_seconds", "tokenizer_types"), + [(1e9, ("openai_tokenizer", "openai_tokenizer")), (1e-9, ("openai_tokenizer", "huggingface_tokenizer"))], + ids=["within_the_window", "after_the_window"], +) +async def test_framed_count_loads_the_tokenizer_again_only_after_the_window_of_a_load_that_fell_back( + monkeypatch: pytest.MonkeyPatch, reuse_seconds: float, tokenizer_types: tuple[str, str] +): + _framed_kimi(monkeypatch) + real_select_tokenizer: Final = litellm.utils._select_tokenizer + hub_tokenizer: Final = litellm.utils.create_tokenizer( + Tokenizer(WordLevel({"[UNK]": 0}, unk_token="[UNK]")).to_str() + ) + loads: Final[list[str]] = [] + + def hub_up_on_the_second_load(model: str, custom_tokenizer: object = None): + loads.append(model) + return real_select_tokenizer(model="gpt-5.4-mini") if len(loads) == 1 else hub_tokenizer + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", hub_up_on_the_second_load) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", reuse_seconds) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 1) + request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]) + + first: Final = await token_counter(request=request, call_endpoint=True) + second: Final = await token_counter(request=request, call_endpoint=True) + + assert (first.tokenizer_type, second.tokenizer_type) == tokenizer_types + + +async def _fork_guard_traceback_frames() -> int: + with pytest.raises(ForkedAfterNativeRuntimeStarted) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + return len(traceback.extract_tb(raised.value.__traceback__)) + + +@pytest.mark.asyncio +async def test_fork_guard_error_while_loading_the_tokenizer_is_raised_to_every_count_sharing_the_load( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi(monkeypatch) + + def forked_load(model: str, custom_tokenizer: object = None) -> NoReturn: + raise ForkedAfterNativeRuntimeStarted("native runtime started before fork") + + monkeypatch.setattr(litellm.utils, "_select_tokenizer", forked_load) + + frames: Final = [await _fork_guard_traceback_frames() for _ in range(3)] + + assert len(set(frames)) == 1 + + +@pytest.mark.asyncio +async def test_swallowed_estimate_error_is_logged_by_type_without_the_request_content( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + verbose_proxy_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name), pytest.raises(ProxyException): + await token_counter( + request=TokenCountRequest( + model="kimi", + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok", "tool_calls": "private-marker-in-the-prompt"}, + ], + ), + call_endpoint=True, + ) + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + warnings: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert any("count error 400" in message and "ValueError" in message for message in warnings) + assert not any("private-marker-in-the-prompt" in message for message in warnings) + + +@pytest.mark.asyncio +async def test_fork_guard_error_in_the_estimate_is_never_answered_as_the_provider_error(): + async def forked_estimate() -> NoReturn: + raise ForkedAfterNativeRuntimeStarted("native runtime started before fork") + + with pytest.raises(ForkedAfterNativeRuntimeStarted): + await _estimate_or_provider_error( + forked_estimate(), provider_failure=ProviderCountFailure(status_code=400, message="scripted 400") + ) + + +@pytest.mark.asyncio +async def test_local_estimate_adds_the_configured_framing_once_per_request(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": f"sk-unused-{name}"}, + "model_info": model_info, + } + + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[deployment("framed", {"token_count_framing": 100}), deployment("plain", {})]), + ) + messages: Final = [{"role": "system", "content": "be brief"}, {"role": "user", "content": "hello there"}] + + plain: Final = await token_counter(request=TokenCountRequest(model="plain", messages=messages)) + framed: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + framed_prompt: Final = await token_counter(request=TokenCountRequest(model="framed", prompt="hello there")) + + assert plain.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + assert framed.total_tokens == plain.total_tokens + 100 + assert (plain.count_source, framed.count_source) == ("local_estimate", "local_estimate") + assert framed_prompt.total_tokens == litellm.token_counter(model="gpt-5.4-mini", text="hello there") + + +@pytest.mark.asyncio +async def test_framing_comes_from_the_row_the_deployment_provider_or_base_model_names(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, litellm_params: dict[str, object], model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {**litellm_params, "api_key": "unused", "api_base": "https://unused.invalid"}, + "model_info": model_info, + } + + kimi: Final = "databricks/system.ai.kimi-k3" + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + deployment("prefixed", {"model": kimi}, {}), + deployment("split", {"model": "system.ai.kimi-k3", "custom_llm_provider": "databricks"}, {}), + deployment("aliased", {"model": "gpt-5.4-mini", "custom_llm_provider": "openai"}, {"base_model": kimi}), + ] + ), + ) + messages: Final = [{"role": "user", "content": "hello there"}] + framing: Final = cost_map_entry(kimi)["token_count_framing"] + + prefixed: Final = await token_counter(request=TokenCountRequest(model="prefixed", messages=messages)) + split: Final = await token_counter(request=TokenCountRequest(model="split", messages=messages)) + aliased: Final = await token_counter(request=TokenCountRequest(model="aliased", messages=messages)) + + assert prefixed.total_tokens == litellm.token_counter(model="system.ai.kimi-k3", messages=messages) + framing + assert split.total_tokens == prefixed.total_tokens + assert aliased.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + framing + + +_KIMI_WEATHER: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +def _byte_level_kimi_k3(path: Path) -> KimiK3TokenizerSource: + path.write_bytes(b"\n".join(base64.b64encode(bytes((value,))) + b" %d" % value for value in range(256))) + return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + + +def _available_kimi_k3(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> tiktoken.Encoding: + source: Final = _byte_level_kimi_k3(tmp_path / "tiktoken.model") + monkeypatch.setattr( + litellm.proxy.proxy_server, "_KIMI_K3_ENCODINGS", KimiK3Encodings(source=source, retry_after_seconds=60) + ) + encoding: Final = load_kimi_k3_encoding(source) + assert encoding is not None + return encoding + + +def _shipped_kimi_k3_framing() -> int: + framing: Final = cost_map_entry("bedrock/us.moonshotai.kimi-k3")["token_count_framing"] + assert isinstance(framing, int) + return framing + + +@pytest.mark.asyncio +async def test_bedrock_kimi_k3_is_counted_with_its_tokenizer_plus_the_shipped_framing_in_both_request_formats( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + _framed_kimi(monkeypatch) + encoding: Final = _available_kimi_k3(monkeypatch, tmp_path) + question: Final = {"role": "user", "content": "Weather in Paris?"} + openai_messages: Final = [ + {"role": "system", "content": "Be terse."}, + question, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}} + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ] + + openai: Final = await count_request_tokens( + request=TokenCountRequest(model="kimi", messages=openai_messages, tools=[_KIMI_WEATHER]), call_endpoint=True + ) + anthropic: Final = await count_request_tokens( + request=TokenCountRequest( + model="kimi", + system="Be terse.", + messages=[ + question, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "c1", "name": "get_weather", "input": {"city": "Paris"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "Sun"}]}, + ], + tools=[ + { + "name": "get_weather", + "description": "Get the weather", + "input_schema": _KIMI_WEATHER["function"]["parameters"], + } + ], + ), + call_endpoint=True, + messages_format="anthropic", + ) + exact: Final = kimi_k3_tokens(encoding, "us.moonshotai.kimi-k3", openai_messages, [_KIMI_WEATHER]) + + assert exact is not None + assert (openai.total_tokens, openai.tokenizer_type, openai.count_source) == ( + exact + _shipped_kimi_k3_framing(), + "huggingface_tokenizer", + "local_estimate", + ) + assert (anthropic.total_tokens, anthropic.tokenizer_type) == (openai.total_tokens, "huggingface_tokenizer") + + +@pytest.mark.asyncio +async def test_bedrock_kimi_k3_deployment_with_a_custom_tokenizer_keeps_counting_with_it( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/kimi-custom-tokenizer") + _available_kimi_k3(monkeypatch, tmp_path) + custom: Final = litellm.utils._select_tokenizer(model="gpt-5.4-mini") + monkeypatch.setattr(litellm.utils, "_select_tokenizer", lambda model, custom_tokenizer=None: custom) + messages: Final = [{"role": "user", "content": "What is new?"}] + + counted: Final = await token_counter(request=TokenCountRequest(model="kimi", messages=messages), call_endpoint=True) + + assert counted.total_tokens == ( + litellm.token_counter(model=counted.model_used, messages=messages, custom_tokenizer=custom) + + kimi_k3_template_tokens(len(messages), False) + + _shipped_kimi_k3_framing() + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("messages", "tools", "template_tokens"), + [ + pytest.param([{"role": "user", "content": "What is new?"}], None, 14, id="one_message"), + pytest.param( + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": "Sunny."}, + {"role": "user", "content": "And Oslo?"}, + ], + [_KIMI_WEATHER], + 14 + 3 * 13 + 40, + id="turns_and_tools", + ), + ], +) +async def test_bedrock_kimi_k3_without_its_tokenizer_adds_the_template_offsets_to_the_local_estimate( + monkeypatch: pytest.MonkeyPatch, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + template_tokens: int, +): + _framed_kimi(monkeypatch) + + counted: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=messages, tools=tools), call_endpoint=True + ) + estimate: Final = litellm.token_counter( + model=counted.model_used, + messages=messages, + tools=tools, + custom_tokenizer=litellm.utils._select_tokenizer(model=counted.model_used), + ) + + assert counted.total_tokens == estimate + template_tokens + _shipped_kimi_k3_framing() + + +@pytest.mark.asyncio +async def test_slow_cost_map_lookup_does_not_stall_other_requests_on_the_worker(monkeypatch: pytest.MonkeyPatch): + real_get_model_info: Final = litellm.get_model_info + lookup_started: Final = threading.Event() + other_request_ran: Final = threading.Event() + released_by_the_other_request: Final[list[bool]] = [] + + def slow_get_model_info(model: str, *args: object, **kwargs: object): + if model == "slow-lookup-model": + lookup_started.set() + released_by_the_other_request.append(other_request_ran.wait(5)) + return real_get_model_info(model, *args, **kwargs) + + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "get_model_info", slow_get_model_info) + counting: Final = asyncio.create_task( + token_counter( + request=TokenCountRequest(model="slow-lookup-model", messages=[{"role": "user", "content": "hi"}]) + ) + ) + await asyncio.to_thread(lookup_started.wait, 5) + other_request_ran.set() + await counting + + assert set(released_by_the_other_request) == {True} + + +@pytest.mark.asyncio +async def test_uncountable_model_keeps_the_provider_error_when_local_counting_is_disabled( + monkeypatch: pytest.MonkeyPatch, +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "kimi", + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": {"token_count_framing": 78}, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + monkeypatch.setattr(litellm, "disable_token_counter", True) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert raised.value.code == "400" + assert "scripted 400 for kimi" in raised.value.message + + +@pytest.mark.parametrize( + ("model_info", "counted_with"), + [ + ({}, ("us.anthropic.claude-opus-5-5", "us-east-1", 447)), + ( + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}}, + ("us.anthropic.claude-opus-5-5", "ap-southeast-4", 447), + ), + ({"token_count_model": {}}, ("us.anthropic.claude-sonnet-5-5", "us-west-2", 300)), + ], +) +@pytest.mark.usefixtures("bedrock_sonnet_5_5_mapped_as_on_main") +@pytest.mark.asyncio +async def test_configured_token_count_model_overrides_the_shipped_one_and_an_empty_one_counts_the_model_itself( + monkeypatch: pytest.MonkeyPatch, model_info: dict[str, object], counted_with: tuple[str, str, int] +): + counter: Final = _ScriptedCounter( + { + "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=300), + "us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447), + } + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-sonnet-5-5", + "aws_region_name": "us-west-2", + }, + "model_info": model_info, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + counted: Final = await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert (counter.calls[0][0], counter.regions[0], counted.total_tokens) == counted_with + assert len(counter.calls) == 1 + + +@pytest.mark.parametrize( + ("aws_region_name", "counted_in"), + [("eu-central-1", "eu-central-1"), ("us-west-2", "us-east-1")], +) +@pytest.mark.usefixtures("bedrock_sonnet_5_5_mapped_as_on_main") +@pytest.mark.asyncio +async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geography( + monkeypatch: pytest.MonkeyPatch, aws_region_name: str, counted_in: str +): + counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": { + "model": "bedrock/anthropic.claude-sonnet-5-5", + "aws_region_name": aws_region_name, + }, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert (counter.calls[0][0], counter.regions[0]) == ("anthropic.claude-opus-5-5", counted_in) + + +class _SlowVault: + def sync_read_secret(self, secret_name: str) -> None: + time.sleep(0.5) + + +async def _longest_stall(done: asyncio.Event) -> float: + if done.is_set(): + return 0.0 + started: Final = time.monotonic() + await asyncio.sleep(0.01) + return max(time.monotonic() - started, await _longest_stall(done)) + + +@pytest.mark.usefixtures("bedrock_sonnet_5_5_mapped_as_on_main") +@pytest.mark.asyncio +async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_free(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + + counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + + def bedrock_counter(deployment: object, model: str) -> tuple[_ScriptedCounter, str, str]: + return counter, model, "bedrock" + + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[{"model_name": "sonnet", "litellm_params": {"model": "bedrock/anthropic.claude-sonnet-5-5"}}] + ), + ) + monkeypatch.setattr(proxy_server, "_get_provider_token_counter", bedrock_counter) + monkeypatch.setattr(litellm, "secret_manager_client", _SlowVault()) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.HASHICORP_VAULT) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + done: Final = asyncio.Event() + + async def count() -> TokenCountResponse: + try: + return await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), + call_endpoint=True, + ) + finally: + done.set() + + started: Final = time.monotonic() + stall, response = await asyncio.gather(_longest_stall(done), count()) + + assert time.monotonic() - started >= 1.0 + assert stall < 0.4 + assert response.total_tokens == 447 + + +@pytest.mark.asyncio +async def test_framing_still_applies_while_the_deployment_is_cooling_down(monkeypatch: pytest.MonkeyPatch): + router: Final = Router( + model_list=[ + { + "model_name": "framed", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-unused-framed"}, + "model_info": {"id": "framed-deployment", "token_count_framing": 100}, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + messages: Final = [{"role": "user", "content": "hello there"}] + warm: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + + router.cooldown_cache.add_deployment_to_cooldown( + model_id="framed-deployment", + original_exception=Exception("rate limited"), + exception_status=429, + cooldown_time=60, + ) + cooling: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + + assert warm.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + 100 + assert cooling.total_tokens - 100 == litellm.token_counter(model="framed", messages=messages) + assert cooling.count_source == "local_estimate" + + +class _GeminiPart(BaseModel): + model_config = ConfigDict(frozen=True) + text: str = "" + + +class _GeminiContent(BaseModel): + model_config = ConfigDict(frozen=True) + parts: tuple[_GeminiPart, ...] = () + + +class _GeminiCountTokensBody(BaseModel): + model_config = ConfigDict(frozen=True) + contents: tuple[_GeminiContent, ...] = () + + +def _word_count(content: _GeminiContent) -> int: + return sum(len(part.text.split()) for part in content.parts) + + +def _gemini_word_counting_peer(request: httpx.Request) -> httpx.Response: + body: Final = _GeminiCountTokensBody.model_validate_json(request.content) + total: Final = sum(_word_count(content) for content in body.contents) + return httpx.Response(200, json={"totalTokens": total}, request=request) + + +def _unavailable_peer(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"error": {"message": "unavailable"}}, request=request) + + +def _gemini_router(drop_params: bool = False) -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-counted", + "litellm_params": { + "model": "gemini/gemini-3.8-flash", + "api_key": "test-key", + "api_base": "https://count-tokens.test", + "drop_params": drop_params, + }, + } + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "text", ["hello there", "The quick brown fox jumps over the lazy dog. " * 50], ids=["short", "long"] +) +async def test_call_endpoint_counts_a_prompt_like_the_same_text_sent_as_a_user_message( + text: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=text), call_endpoint=True + ) + from_message: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": text}]), + call_endpoint=True, + ) + + assert from_prompt.error is False + assert from_prompt.total_tokens == from_message.total_tokens == len(text.split()) + + +class _PartnerMessage(BaseModel): + model_config = ConfigDict(frozen=True) + content: str = "" + + +class _PartnerCountRequest(BaseModel): + model_config = ConfigDict(frozen=True) + messages: tuple[_PartnerMessage, ...] = () + + +async def _vertex_partner_word_count( + model: str, request_data: JsonValue, litellm_params: JsonValue +) -> dict[str, JsonValue]: + body: Final = _PartnerCountRequest.model_validate(request_data) + return {"input_tokens": sum(len(message.content.split()) for message in body.messages)} + + +@pytest.mark.asyncio +async def test_call_endpoint_counts_a_prompt_on_a_vertex_ai_deployment(monkeypatch: pytest.MonkeyPatch) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "vertex-counted", + "litellm_params": { + "model": "vertex_ai/claude-sonnet-5-5", + "vertex_project": "test-project", + "vertex_location": "us-east5", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + text: Final = "The quick brown fox jumps over the lazy dog." + + with patch.object( + VertexAIPartnerModelsTokenCounter, "handle_count_tokens_request", side_effect=_vertex_partner_word_count + ): + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="vertex-counted", prompt=text), call_endpoint=True + ) + + assert (from_prompt.error, from_prompt.total_tokens) == (False, 9) + + +_FOX: Final = "The quick brown fox jumps over the lazy dog." + + +_CONVERSATIONS: Final = ( + {"messages": [{"role": "user", "content": _FOX}]}, + {"contents": [{"role": "user", "parts": [{"text": _FOX}]}]}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("conversation", _CONVERSATIONS, ids=["messages", "contents"]) +@pytest.mark.parametrize("call_endpoint", [True, False], ids=["provider", "local"]) +async def test_prompt_next_to_a_conversation_answers_400_without_drop_params( + conversation: dict[str, JsonValue], call_endpoint: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with ( + patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), + pytest.raises(ProxyException) as rejected, + ): + await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint + ) + + assert (rejected.value.code, rejected.value.param) == ("400", "prompt") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("conversation", "call_endpoint", "global_drop_params", "deployment_drop_params"), + [ + (_CONVERSATIONS[0], True, False, True), + (_CONVERSATIONS[1], True, False, True), + (_CONVERSATIONS[0], False, False, True), + (_CONVERSATIONS[0], False, True, False), + ], + ids=[ + "provider_messages", + "provider_contents", + "local_messages", + "local_messages_global_setting", + ], +) +async def test_drop_params_counts_a_conversation_without_the_prompt( + conversation: dict[str, JsonValue], + call_endpoint: bool, + global_drop_params: bool, + deployment_drop_params: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router(drop_params=deployment_drop_params)) + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + with_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint + ) + conversation_only: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", **conversation), call_endpoint=call_endpoint + ) + + assert with_prompt.error is False + assert with_prompt.total_tokens == conversation_only.total_tokens > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("empty_conversation", [{"messages": []}, {"contents": []}], ids=["messages", "contents"]) +async def test_call_endpoint_prompt_next_to_an_empty_conversation_counts_the_prompt( + empty_conversation: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + counted: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation), call_endpoint=True + ) + + assert (counted.error, counted.total_tokens) == (False, 9) + + +@pytest.mark.asyncio +async def test_local_count_of_a_prompt_next_to_empty_messages_counts_the_prompt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + + with_empty: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX, messages=[])) + alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) + + assert with_empty.error is False + assert with_empty.total_tokens == alone.total_tokens > 0 + + +_WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chat_fields", + [{"system": "You answer in one short sentence. " * 20}, {"tools": [_WEATHER_TOOL]}], + ids=["system", "tools"], +) +@pytest.mark.parametrize("empty_conversation", [{}, {"messages": []}], ids=["no_messages", "empty_messages"]) +async def test_local_count_of_a_prompt_with_system_or_tools_counts_them_like_the_message_form( + chat_fields: dict[str, JsonValue], empty_conversation: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation, **chat_fields) + ) + from_message: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": _FOX}], **chat_fields) + ) + prompt_alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) + + assert from_prompt.error is False + assert from_prompt.total_tokens == from_message.total_tokens > prompt_alone.total_tokens + + +@pytest.mark.asyncio +async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_needs_messages( + monkeypatch: pytest.MonkeyPatch, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "anthropic-counted", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_key": "test-key", + "api_base": "https://count-tokens.test", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_unavailable_peer)) + request: Final = TokenCountRequest(model="anthropic-counted", prompt="The quick brown fox jumps over the lazy dog.") + + with patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client): + with_endpoint: Final = await token_counter(request=request, call_endpoint=True) + local: Final = await token_counter(request=request, call_endpoint=False) + + assert (with_endpoint.error, with_endpoint.total_tokens, with_endpoint.tokenizer_type) == ( + False, + local.total_tokens, + local.tokenizer_type, + ) diff --git a/tests/unit/test_count_tokens_public_api.py b/tests/unit/test_count_tokens_public_api.py index 84121dd63e4..e33bc1e5907 100644 --- a/tests/unit/test_count_tokens_public_api.py +++ b/tests/unit/test_count_tokens_public_api.py @@ -97,7 +97,7 @@ def test_acount_tokens_with_tools(): assert result.total_tokens == 30 mock_handler.assert_called_once() call_kwargs = mock_handler.call_args - assert call_kwargs.kwargs.get("tools") == tools + assert list(call_kwargs.kwargs.get("tools") or ()) == tools def test_acount_tokens_with_system(): diff --git a/tests/unit/test_model_prices_schema.py b/tests/unit/test_model_prices_schema.py index d60e9e85f45..766dc9d8648 100644 --- a/tests/unit/test_model_prices_schema.py +++ b/tests/unit/test_model_prices_schema.py @@ -176,6 +176,52 @@ def test_generated_off_peak_schema_rejects_malformed_blocks(block: dict): assert not validator.is_valid({"some-model": {**OFF_PEAK_ENTRY, "off_peak_pricing": block}}) +TOKEN_COUNT_ENTRY: Final = MappingProxyType( + { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "token_count_framing": 78, + "token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"}, + } +) + + +@pytest.mark.parametrize( + "setting", + [ + {"token_count_framing": 0}, + {"token_count_framing": 78}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5"}}, + ], +) +def test_generator_classifies_token_count_settings(setting: dict): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) + assert build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}}) + + +@pytest.mark.parametrize( + "setting", + [ + {"token_count_framing": -1}, + {"token_count_framing": "78"}, + {"token_count_framing": 1.5}, + {"token_count_model": {"model": ""}}, + {"token_count_model": {"model": " "}}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": ""}}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": " "}}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": 5}}, + {"token_count_model": {"region": "us-east-1"}}, + {"token_count_model": "us.anthropic.claude-opus-5-5"}, + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "zone": "us-east-1"}}, + ], +) +def test_generated_token_count_schema_rejects_malformed_settings(setting: dict): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) + assert not build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}}) + + def find_duplicate_keys(path: Path) -> list[str]: duplicates: list[str] = [] diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index ecc1f3ab078..7a865b8df9b 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -548,6 +548,40 @@ def test_regex_lookaround_flag_stays_on_the_deployment_that_set_it() -> None: _restore_model_cost_entries(model_keys) +def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None: + """A deployment's ``token_count_model`` and ``token_count_framing`` must not land on + the shared ``{provider}/{model}`` key, or a sibling deployment of that model would count with them.""" + from litellm.proxy.common_utils.token_count_settings import cost_map_entry + + backend_model: Final = "bedrock/us.moonshotai.kimi-k3" + deploy_id: Final = "kimi-deploy-keep-token-count" + configured: Final = { + "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-west-2"}, + "token_count_framing": 7, + } + builtin: Final = {field: cost_map_entry(backend_model).get(field) for field in configured} + assert all(builtin[field] != value for field, value in configured.items()), builtin + model_keys: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in (deploy_id, backend_model, backend_model.removeprefix("bedrock/")) + } + try: + Router( + model_list=[ + { + "model_name": "kimi-keep-token-count", + "litellm_params": {"model": backend_model}, + "model_info": {"id": deploy_id, **configured}, + } + ], + ) + + assert {field: litellm.model_cost[deploy_id][field] for field in configured} == configured + assert {field: cost_map_entry(backend_model).get(field) for field in configured} == builtin + finally: + _restore_model_cost_entries(model_keys) + + def test_should_store_full_pricing_under_deployment_model_id(): """ Per-deployment pricing (including zero) should be stored and diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 80d97d3e99b..40a676fd113 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1042,6 +1042,13 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_bedrock_runtime_chat_completions_tools_with_reasoning": {"type": "boolean"}, "supports_bedrock_runtime_chat_completions_response_format": {"type": "boolean"}, "supports_bedrock_runtime_chat_completions_inline_reasoning": {"type": "boolean"}, + "token_count_framing": {"type": "integer"}, + "token_count_model": { + "type": "object", + "properties": {"model": {"type": "string"}, "region": {"type": "string"}}, + "required": ["model"], + }, + "token_count_tokenizer": {"type": "string", "enum": ["kimi_k3"]}, "supports_url_context": {"type": "boolean"}, "supports_multimodal": {"type": "boolean"}, "uses_embed_content": {"type": "boolean"}, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index df25c1b1c67..1fdf5776b1f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -48025,9 +48025,17 @@ export interface components { * @description Response structure for token count details with modality breakdown. * * Example: - * {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]} + * {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}], + * 'litellm_count_source': 'provider'} */ TokenCountDetailsResponse: { + /** Litellm Count Model */ + litellm_count_model?: string; + /** + * Litellm Count Source + * @enum {string} + */ + litellm_count_source?: "provider" | "token_count_model" | "local_estimate"; /** Prompttokensdetails */ promptTokensDetails: components["schemas"]["PromptTokensDetails"][]; /** Totaltokens */ @@ -48056,6 +48064,8 @@ export interface components { }; /** TokenCountResponse */ TokenCountResponse: { + /** Count Source */ + count_source?: ("provider" | "token_count_model" | "local_estimate") | null; /** * Error * @default false From 98d702fba93958d501b8a93661d1273bac1a2185 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:29:29 +0000 Subject: [PATCH 08/13] Merge pull request #47446 from BerriAI/litellm_backport_google_count_shapes_stable_1_104_x fix(count_tokens): restore Google request compatibility on stable/1.104.x (cherry picked from commit a63df167da30438b08fb1568b340ab787a4e182a) --- litellm/proxy/google_endpoints/endpoints.py | 11 ++++-- .../test_google_gemini_proxy_request.py | 39 ++++++++++++++++++- 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index e15daf05ced..40494597a68 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -5,7 +5,7 @@ from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse, ORJSONResponse -from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, field_validator from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types from litellm.proxy._types import * @@ -34,6 +34,11 @@ class _GoogleCountTokensSystemInstruction(BaseModel): model_config = ConfigDict(frozen=True) parts: tuple[_GoogleCountTokensTextPart, ...] + @field_validator("parts", mode="before") + @classmethod + def normalize_parts(cls, parts: object) -> object: + return (parts,) if isinstance(parts, Mapping) else parts + class _GoogleCountTokensFunctionDeclaration(BaseModel): model_config = ConfigDict(frozen=True, populate_by_name=True) @@ -54,7 +59,7 @@ class _GoogleCountTokensRequest(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True) contents: tuple[Mapping[str, JsonValue], ...] = () system_instruction: _GoogleCountTokensSystemInstruction | None = Field(default=None, alias="systemInstruction") - tools: tuple[_GoogleCountTokensTool, ...] = () + tools: tuple[_GoogleCountTokensTool, ...] | None = () router: Final = APIRouter( @@ -114,7 +119,7 @@ def _google_count_tokens_request(model_name: str, data: object) -> TokenCountReq if validated_data.system_instruction is not None else None ) - declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools) + declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools or ()) tools: Final = [ { "name": declaration.name, diff --git a/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py b/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py index 6f8f90efc73..c465b6167e6 100644 --- a/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py +++ b/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py @@ -9,11 +9,12 @@ The request payload is correctly processed and forwarded to the httpx client. import json import unittest.mock -from typing import Optional +from typing import Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from pydantic import JsonValue, ValidationError # Add the parent directory to the system path @@ -493,3 +494,39 @@ async def test_gemini_proxy_config_with_custom_api_base(): if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +@pytest.mark.parametrize("system_field", ("systemInstruction", "system_instruction")) +@pytest.mark.parametrize("parts", ({"text": "Follow these rules"}, [{"text": "Follow these rules"}])) +@pytest.mark.parametrize("tools", (None, [])) +def test_google_count_tokens_accepts_equivalent_system_parts_and_no_tools( + system_field: str, parts: JsonValue, tools: JsonValue +) -> None: + from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request + + contents: Final = [{"role": "user", "parts": [{"text": "Count this request"}]}] + request: Final = _google_count_tokens_request( + "counted", {"contents": contents, system_field: {"parts": parts}, "tools": tools} + ) + + assert request.messages == [{"role": "user", "content": "Count this request"}] + assert request.system == [{"type": "text", "text": "Follow these rules"}] + assert request.tools is None + assert request.contents == contents + + +@pytest.mark.parametrize( + "parts", + ( + {"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}, + [{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}], + ), +) +def test_google_count_tokens_rejects_media_only_system_parts(parts: JsonValue) -> None: + from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request + + with pytest.raises(ValidationError): + _google_count_tokens_request( + "counted", + {"contents": [{"role": "user", "parts": [{"text": "Hello"}]}], "systemInstruction": {"parts": parts}}, + ) From 699e79c04ea7bf6c25c38987023dd38ca083f154 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Mon, 19 Oct 2026 17:59:47 -0700 Subject: [PATCH 09/13] Merge pull request #47454 from BerriAI/litellm_restore_token_count_input_precedence fix(proxy): restore legacy token count input precedence (cherry picked from commit 5346a40c37a25b6e264c7b3783c5afb7d99ba456) --- litellm/proxy/proxy_server.py | 12 ------ .../test_vertex_gemini_count_tokens_wire.py | 38 +++++++++++++++- tests/unit/proxy/test_proxy_token_counter.py | 43 ++++--------------- 3 files changed, 45 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 97001ad7b30..77203399411 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14588,11 +14588,6 @@ _PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( ) -def _drops_params(deployment: Mapping[str, object] | None) -> bool: - litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) - return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) - - def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -14680,13 +14675,6 @@ async def count_request_tokens( model_to_use = _model conversation_empty: Final = not messages and not contents - if prompt and not conversation_empty and not _drops_params(deployment): - raise ProxyException( - message="prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt", - type="invalid_request_error", - param="prompt", - code=400, - ) prompt_only: Final = bool(prompt) and conversation_empty prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES ladder: Final = ( diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 5fe82d37e9e..e54fb1a4a13 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -220,8 +220,8 @@ _CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingPro ) -def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str: - return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key") +def _ai_studio_deployment(scenario: Scenario, api_base: str, **parameters: JsonValue) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key", **parameters) def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: @@ -594,3 +594,37 @@ def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gatew tuple(_json_body(request) for request in requests) == ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9 ) + + +@pytest.mark.parametrize( + ("conversation", "call_endpoint"), + ( + ({"messages": _MESSAGES}, True), + ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]}, True), + ({"messages": _MESSAGES}, False), + ), + ids=("provider-count-messages", "provider-count-contents", "local-count-messages"), +) +@pytest.mark.parametrize("drop_params", (False, True), ids=("no-drop", "drop")) +def test_utils_token_counter_counts_the_conversation_regardless_of_prompt_and_drop_params( + gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool, drop_params: bool +) -> None: + peer: Final = _GeminiPeer() + params: Final = {"call_endpoint": "true" if call_endpoint else "false"} + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=drop_params) + mixed: Final = gateway.request( + "POST", "/utils/token_counter", {"model": model, "prompt": "hi", **conversation}, params=params + ) + mixed_requests: Final = _count_requests(wire) + alone: Final = gateway.request("POST", "/utils/token_counter", {"model": model, **conversation}, params=params) + alone_requests: Final = _count_requests(wire) + assert mixed.status_code == 200, mixed.text + assert alone.status_code == 200, alone.text + mixed_tokens: Final = TypeAdapter(int).validate_python(_JSON_OBJECT.validate_json(mixed.content)["total_tokens"]) + alone_tokens: Final = TypeAdapter(int).validate_python(_JSON_OBJECT.validate_json(alone.content)["total_tokens"]) + assert mixed_tokens == alone_tokens > 0 + assert len(mixed_requests) == len(alone_requests) == int(call_endpoint) + assert [_counted_contents(_json_body(request)) for request in mixed_requests] == [ + _counted_contents(_json_body(request)) for request in alone_requests + ] diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 4e0b242f80b..8b0b55f0af0 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2511,43 +2511,18 @@ _CONVERSATIONS: Final = ( ) -@pytest.mark.asyncio -@pytest.mark.parametrize("conversation", _CONVERSATIONS, ids=["messages", "contents"]) -@pytest.mark.parametrize("call_endpoint", [True, False], ids=["provider", "local"]) -async def test_prompt_next_to_a_conversation_answers_400_without_drop_params( - conversation: dict[str, JsonValue], call_endpoint: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) - client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) - - with ( - patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), - pytest.raises(ProxyException) as rejected, - ): - await token_counter( - request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint - ) - - assert (rejected.value.code, rejected.value.param) == ("400", "prompt") - - @pytest.mark.asyncio @pytest.mark.parametrize( - ("conversation", "call_endpoint", "global_drop_params", "deployment_drop_params"), - [ - (_CONVERSATIONS[0], True, False, True), - (_CONVERSATIONS[1], True, False, True), - (_CONVERSATIONS[0], False, False, True), - (_CONVERSATIONS[0], False, True, False), - ], - ids=[ - "provider_messages", - "provider_contents", - "local_messages", - "local_messages_global_setting", - ], + ("conversation", "call_endpoint"), + ((_CONVERSATIONS[0], True), (_CONVERSATIONS[1], True), (_CONVERSATIONS[0], False)), + ids=("provider_messages", "provider_contents", "local_messages"), ) -async def test_drop_params_counts_a_conversation_without_the_prompt( +@pytest.mark.parametrize( + ("global_drop_params", "deployment_drop_params"), + ((False, False), (False, True), (True, False), (True, True)), + ids=("no_drop", "deployment_drop", "global_drop", "both_drop"), +) +async def test_conversation_takes_precedence_over_prompt_regardless_of_drop_params( conversation: dict[str, JsonValue], call_endpoint: bool, global_drop_params: bool, From 1fe5581b27f2529158dec339e3e800a230da6b81 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 19 Oct 2026 23:09:01 -0700 Subject: [PATCH 10/13] chore: bump litellm 1.106.0 -> 1.106.1 litellm 1.106.0 already shipped, so the next release off this line is 1.106.1 --- pyproject.toml | 4 ++-- uv.lock | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e59e452650b..aa7e597fe44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.106.0" +version = "1.106.1" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -373,7 +373,7 @@ litellm-enterprise = { workspace = true } members = ["enterprise", "litellm-proxy-extras"] [tool.commitizen] -version = "1.106.0" +version = "1.106.1" version_files = [ "pyproject.toml:^version", ] diff --git a/uv.lock b/uv.lock index 96fa27289fa..52713523e0a 100644 --- a/uv.lock +++ b/uv.lock @@ -4500,7 +4500,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.106.0" +version = "1.106.1" source = { editable = "." } dependencies = [ { name = "aiohttp" }, From e46904b54602cfe98c9a43c0c1442f35e6c34ed2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 20 Oct 2026 01:20:54 -0700 Subject: [PATCH 11/13] fix(anthropic): carry only thinking.block_binding through the legacy thinking rebuild The #46342 pick spread every caller key from a legacy thinking.type=enabled request into the rebuilt adaptive param, so block_binding and any unknown key newly reached Bedrock Converse, Vertex AI and Databricks, where 1.106.0 sent only {"type": "adaptive"}. The rebuild now carries block_binding alone, and only for Anthropic, Azure AI and Bedrock Invoke. The Bedrock Invoke display: updates and clear_thinking rebuilds keep block_binding and drop the other caller keys, as 1.106.0 did --- litellm/llms/anthropic/chat/transformation.py | 10 +- litellm/llms/anthropic/common_utils.py | 13 ++- .../pass_through/messages/transformation.py | 6 ++ .../anthropic_claude3_transformation.py | 17 +++- .../anthropic_claude3_transformation.py | 24 ++++- .../anthropic/test_anthropic_common_utils.py | 91 +++++++++++++++---- ...ations_anthropic_claude3_transformation.py | 10 +- .../test_anthropic_claude3_transformation.py | 6 +- 8 files changed, 141 insertions(+), 36 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 54a8780ca5b..1da94f59668 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -97,6 +97,7 @@ from litellm.utils import ( ) from ..common_utils import ( + THINKING_BLOCK_BINDING_PROVIDERS, AnthropicError, AnthropicModelInfo, eager_input_streaming_flag, @@ -329,6 +330,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _resolved_provider(self) -> str: return self.custom_llm_provider or "anthropic" + @property + def _carries_thinking_block_binding(self) -> bool: + return self._resolved_provider in THINKING_BLOCK_BINDING_PROVIDERS + @classmethod def get_config(cls, *, model: str | None = None): config: Final = super().get_config() @@ -1656,7 +1661,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): else: optional_params["thinking"] = value AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( - model=model, optional_params=optional_params, custom_llm_provider=self._resolved_provider + model=model, + optional_params=optional_params, + custom_llm_provider=self._resolved_provider, + carry_block_binding=self._carries_thinking_block_binding, ) elif param == "reasoning_effort": # Accept both string ("low") and dict ({"effort": "low", diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 12d77112b08..2e184ee82fe 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -108,6 +108,7 @@ _DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)") _CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:" _OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) _CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +THINKING_BLOCK_BINDING_PROVIDERS: Final = frozenset({"anthropic", "azure_ai"}) _CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/") @@ -162,6 +163,11 @@ def validated_object_mapping(value: object) -> dict[object, object] | None: return None +def thinking_block_binding(thinking: Mapping[object, object] | None) -> dict[str, object]: + block_binding: Final = thinking.get("block_binding") if thinking is not None else None + return {} if block_binding is None else {"block_binding": block_binding} + + def _validated_claude_code_list(value: object) -> list[object] | None: try: return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value) @@ -914,6 +920,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): model: str, optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param like the sibling helpers custom_llm_provider: str, + carry_block_binding: bool = False, ) -> None: """Translate legacy ``thinking.type=enabled`` to adaptive for the adaptive-thinking models that reject it (4.7+ and the 5 families). @@ -939,10 +946,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): existing_output_config: Final = optional_params.get("output_config") display: Final = thinking.get("display") supported_display: Final = {"display": display} if display in ("summarized", "omitted") else {} - carried_thinking_keys: Final = { - key: value for key, value in thinking.items() if key not in ("type", "budget_tokens", "display") - } - optional_params["thinking"] = {**carried_thinking_keys, "type": "adaptive", **supported_display} + carried_block_binding: Final = thinking_block_binding(thinking) if carry_block_binding else {} + optional_params["thinking"] = {**carried_block_binding, "type": "adaptive", **supported_display} optional_params["output_config"] = { "effort": effort, **(existing_output_config if isinstance(existing_output_config, dict) else MappingProxyType({})), diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index f1eb30b63c1..9558a0012a5 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -23,6 +23,7 @@ from litellm.types.llms.anthropic_tool_search import get_tool_search_beta_header from litellm.types.router import GenericLiteLLMParams from ...common_utils import ( + THINKING_BLOCK_BINDING_PROVIDERS, AnthropicError, AnthropicModelInfo, merge_anthropic_beta_headers, @@ -80,6 +81,10 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): def _resolved_provider(self) -> str: return self.custom_llm_provider or "anthropic" + @property + def _carries_thinking_block_binding(self) -> bool: + return self._resolved_provider in THINKING_BLOCK_BINDING_PROVIDERS + def get_supported_anthropic_messages_params(self, model: str) -> list: return [ "messages", @@ -607,6 +612,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): model=model, optional_params=anthropic_messages_optional_request_params, custom_llm_provider=self._resolved_provider, + carry_block_binding=self._carries_thinking_block_binding, ) self._translate_adaptive_effort_for_non_adaptive_model( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index b8bc276de0d..24c219e9837 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -12,7 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( convert_url_to_base64, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig -from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.anthropic.common_utils import AnthropicModelInfo, thinking_block_binding from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, ) @@ -67,6 +67,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): def custom_llm_provider(self) -> str | None: return "bedrock" + @property + def _carries_thinking_block_binding(self) -> bool: + return True + def should_strip_billing_metadata(self) -> bool: return True @@ -115,7 +119,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): model = original_model AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( - model=original_model, optional_params=optional_params, custom_llm_provider="bedrock" + model=original_model, + optional_params=optional_params, + custom_llm_provider="bedrock", + carry_block_binding=self._carries_thinking_block_binding, ) translated_thinking: Final = optional_params.get("thinking") if ( @@ -123,7 +130,11 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): and isinstance(translated_thinking, dict) and translated_thinking.get("type") == "adaptive" ): - optional_params["thinking"] = {**translated_thinking, "display": "updates"} + optional_params["thinking"] = { + **thinking_block_binding(translated_thinking), + "type": "adaptive", + "display": "updates", + } # The stub model hides the original model from the parent's forced-tool-use backstop response_format_tool_choice: Final = optional_params.get("tool_choice") diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index eb6ec66bb09..1b857e49505 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -17,7 +17,11 @@ from litellm.llms.anthropic.chat.transformation import ( DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, AnthropicConfig, ) -from litellm.llms.anthropic.common_utils import AnthropicModelInfo, validated_object_mapping +from litellm.llms.anthropic.common_utils import ( + AnthropicModelInfo, + thinking_block_binding, + validated_object_mapping, +) from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -81,6 +85,10 @@ class AmazonAnthropicClaudeMessagesConfig( def custom_llm_provider(self) -> str | None: return "bedrock" + @property + def _carries_thinking_block_binding(self) -> bool: + return True + @property def beta_headers_provider(self) -> str: return self.custom_llm_provider or "bedrock" @@ -349,9 +357,13 @@ class AmazonAnthropicClaudeMessagesConfig( output_config = {} output_config.setdefault("effort", self._effort_from_thinking_budget(budget_tokens)) thinking: Final = validated_object_mapping(anthropic_messages_request.get("thinking")) - carried: Final = {key: value for key, value in (thinking or {}).items() if key not in ("type", "budget_tokens")} + display: Final = thinking.get("display") if thinking is not None else None anthropic_messages_request["output_config"] = output_config - anthropic_messages_request["thinking"] = {**carried, "type": "adaptive"} + anthropic_messages_request["thinking"] = { + **thinking_block_binding(thinking), + "type": "adaptive", + **({} if display is None else {"display": display}), + } verbose_logger.debug( "Bedrock clear_thinking_20251015: injected adaptive thinking with effort=%s for model=%s", output_config.get("effort"), @@ -692,7 +704,11 @@ class AmazonAnthropicClaudeMessagesConfig( and isinstance(translated_thinking, dict) and translated_thinking.get("type") == "adaptive" ): - anthropic_messages_request["thinking"] = {**translated_thinking, "display": "updates"} + anthropic_messages_request["thinking"] = { + **thinking_block_binding(translated_thinking), + "type": "adaptive", + "display": "updates", + } self._normalize_system_role_messages(anthropic_messages_request, model=model) ######################################################### ############## BEDROCK Invoke SPECIFIC TRANSFORMATION ### diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index dcf0d58d675..05aa3bfe5ef 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -16,7 +16,7 @@ import sys import threading from collections.abc import Mapping from pathlib import Path -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import Final from unittest.mock import patch @@ -4172,28 +4172,81 @@ def test_shared_legacy_thinking_translation_preserves_supported_display( assert optional_params["thinking"] == expected_thinking -@pytest.mark.usefixtures("local_model_cost_map") -def test_shared_legacy_thinking_translation_carries_block_binding() -> None: - from litellm.llms.anthropic.common_utils import AnthropicModelInfo - - optional_params: Final = { - "thinking": { - "type": "enabled", - "budget_tokens": 2048, - "block_binding": {"prefix_mismatch_behavior": "error"}, - }, +LEGACY_THINKING_WITH_BLOCK_BINDING: Final = MappingProxyType( + { + "type": "enabled", + "budget_tokens": 2048, + "block_binding": {"prefix_mismatch_behavior": "error"}, + "future_key": 1, } +) +ADAPTIVE_THINKING_WITH_BLOCK_BINDING: Final = MappingProxyType( + {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}} +) - AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( - model="claude-opus-5", - optional_params=optional_params, - custom_llm_provider="azure_ai", + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected_thinking"), + ( + ("claude-opus-4-8", "anthropic", ADAPTIVE_THINKING_WITH_BLOCK_BINDING), + ("claude-opus-4-8", "azure_ai", ADAPTIVE_THINKING_WITH_BLOCK_BINDING), + ("invoke/us.anthropic.claude-opus-4-8", "bedrock", ADAPTIVE_THINKING_WITH_BLOCK_BINDING), + ("converse/us.anthropic.claude-opus-4-8", "bedrock", {"type": "adaptive"}), + ("claude-opus-4-8", "vertex_ai", {"type": "adaptive"}), + ("databricks-claude-opus-4-8", "databricks", {"type": "adaptive"}), + ), + ids=("anthropic", "azure_ai", "bedrock_invoke", "bedrock_converse", "vertex_ai", "databricks"), +) +def test_chat_legacy_thinking_translation_carries_only_block_binding_where_supported( + model: str, custom_llm_provider: str, expected_thinking: Mapping[str, object] +) -> None: + from litellm.utils import get_optional_params + + optional_params: Final = get_optional_params( + model=model, + custom_llm_provider=custom_llm_provider, + thinking=dict(LEGACY_THINKING_WITH_BLOCK_BINDING), ) - assert optional_params["thinking"] == { - "type": "adaptive", - "block_binding": {"prefix_mismatch_behavior": "error"}, - } + assert optional_params["thinking"] == expected_thinking + + +@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected_thinking"), + ( + ("claude-opus-4-8", "anthropic", ADAPTIVE_THINKING_WITH_BLOCK_BINDING), + ("claude-opus-4-8", "azure_ai", ADAPTIVE_THINKING_WITH_BLOCK_BINDING), + ("us.anthropic.claude-opus-4-8", "bedrock", ADAPTIVE_THINKING_WITH_BLOCK_BINDING), + ("claude-opus-4-8", "vertex_ai", {"type": "adaptive"}), + ), + ids=("anthropic", "azure_ai", "bedrock_invoke", "vertex_ai"), +) +def test_messages_legacy_thinking_translation_carries_only_block_binding_where_supported( + model: str, custom_llm_provider: str, expected_thinking: Mapping[str, object] +) -> None: + from litellm.types.router import GenericLiteLLMParams + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + config: Final = ProviderConfigManager.get_provider_anthropic_messages_config( + model=model, provider=LlmProviders(custom_llm_provider) + ) + assert config is not None + + result: Final = config.transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": "Reply with OK"}], + anthropic_messages_optional_request_params={ + "max_tokens": 512, + "thinking": dict(LEGACY_THINKING_WITH_BLOCK_BINDING), + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("thinking") == expected_thinking @pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 1abebaaa8c5..cc154c3d026 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1004,8 +1004,14 @@ BLOCK_BINDING: Final = {"prefix_mismatch_behavior": "error"} @pytest.mark.parametrize( "requested_thinking", ( - {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING}, - {"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING}, + {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, + { + "type": "enabled", + "budget_tokens": 2048, + "display": "updates", + "block_binding": BLOCK_BINDING, + "future_key": 1, + }, ), ids=("adaptive", "legacy"), ) diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 579972f8258..059ef765990 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -3536,15 +3536,15 @@ BLOCK_BINDING: Final = {"prefix_mismatch_behavior": "error"} @pytest.mark.parametrize( ("model", "requested_thinking", "extra_params"), ( - ("eu.anthropic.claude-opus-5", {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING}, {}), + ("eu.anthropic.claude-opus-5", {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, {}), ( "eu.anthropic.claude-opus-5", - {"type": "enabled", "budget_tokens": 24000, "display": "updates", "block_binding": BLOCK_BINDING}, + {"type": "enabled", "budget_tokens": 24000, "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, {}, ), ( "us.anthropic.claude-opus-4-6", - {"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING}, + {"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, {"context_management": {"edits": [{"type": "clear_thinking_20251015"}]}}, ), ), From 1e23ad47917318d6f68b26f7bfd50e560f71b094 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 20 Oct 2026 02:08:52 -0700 Subject: [PATCH 12/13] fix(anthropic): forward thinking.block_binding only from the exact Anthropic, Azure AI and Bedrock Invoke configs e46904b546 decided by resolved provider name, so the OpenAI-compatible /v1/messages passthrough and Bedrock Claude Platform /v1/messages, which inherit the "anthropic" provider name, and the Bedrock Mantle chat and /v1/messages configs, which inherit the Bedrock Invoke override, still sent block_binding where 1.106.0 sent {"type": "adaptive"}. Each in-scope config now opts in for its own class only, and the Bedrock Invoke display: updates and clear_thinking rebuilds keep block_binding only when the config carries it --- litellm/llms/anthropic/chat/transformation.py | 3 +- litellm/llms/anthropic/common_utils.py | 1 - .../pass_through/messages/transformation.py | 3 +- .../anthropic/messages_transformation.py | 4 ++ .../llms/azure_ai/anthropic/transformation.py | 4 ++ .../anthropic_claude3_transformation.py | 4 +- .../anthropic_claude3_transformation.py | 6 +- .../anthropic/test_anthropic_common_utils.py | 51 +++++++++++++--- ...ations_anthropic_claude3_transformation.py | 30 +++++++++- .../test_anthropic_claude3_transformation.py | 59 ++++++++++++++++++- 10 files changed, 143 insertions(+), 22 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 1da94f59668..74ef8936941 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -97,7 +97,6 @@ from litellm.utils import ( ) from ..common_utils import ( - THINKING_BLOCK_BINDING_PROVIDERS, AnthropicError, AnthropicModelInfo, eager_input_streaming_flag, @@ -332,7 +331,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): @property def _carries_thinking_block_binding(self) -> bool: - return self._resolved_provider in THINKING_BLOCK_BINDING_PROVIDERS + return type(self) is AnthropicConfig @classmethod def get_config(cls, *, model: str | None = None): diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 2e184ee82fe..4415089f480 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -108,7 +108,6 @@ _DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)") _CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:" _OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) _CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) -THINKING_BLOCK_BINDING_PROVIDERS: Final = frozenset({"anthropic", "azure_ai"}) _CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/") diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index 9558a0012a5..418be6d95a5 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -23,7 +23,6 @@ from litellm.types.llms.anthropic_tool_search import get_tool_search_beta_header from litellm.types.router import GenericLiteLLMParams from ...common_utils import ( - THINKING_BLOCK_BINDING_PROVIDERS, AnthropicError, AnthropicModelInfo, merge_anthropic_beta_headers, @@ -83,7 +82,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): @property def _carries_thinking_block_binding(self) -> bool: - return self._resolved_provider in THINKING_BLOCK_BINDING_PROVIDERS + return type(self) is AnthropicMessagesConfig def get_supported_anthropic_messages_params(self, model: str) -> list: return [ diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index 6803bf7e551..159ef99f74d 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -22,6 +22,10 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): def custom_llm_provider(self) -> str | None: return "azure_ai" + @property + def _carries_thinking_block_binding(self) -> bool: + return type(self) is AzureAnthropicMessagesConfig + def should_strip_billing_metadata(self) -> bool: return True diff --git a/litellm/llms/azure_ai/anthropic/transformation.py b/litellm/llms/azure_ai/anthropic/transformation.py index 18a4eeb7c0a..ecc5bf9fea4 100644 --- a/litellm/llms/azure_ai/anthropic/transformation.py +++ b/litellm/llms/azure_ai/anthropic/transformation.py @@ -39,6 +39,10 @@ class AzureAnthropicConfig(AnthropicConfig): def custom_llm_provider(self) -> str | None: return "azure_ai" + @property + def _carries_thinking_block_binding(self) -> bool: + return type(self) is AzureAnthropicConfig + def should_strip_billing_metadata(self) -> bool: return True diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 24c219e9837..7f99b647025 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -69,7 +69,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): @property def _carries_thinking_block_binding(self) -> bool: - return True + return type(self) is AmazonAnthropicClaudeConfig def should_strip_billing_metadata(self) -> bool: return True @@ -131,7 +131,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): and translated_thinking.get("type") == "adaptive" ): optional_params["thinking"] = { - **thinking_block_binding(translated_thinking), + **(thinking_block_binding(translated_thinking) if self._carries_thinking_block_binding else {}), "type": "adaptive", "display": "updates", } diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 1b857e49505..d3a3490b372 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -87,7 +87,7 @@ class AmazonAnthropicClaudeMessagesConfig( @property def _carries_thinking_block_binding(self) -> bool: - return True + return type(self) is AmazonAnthropicClaudeMessagesConfig @property def beta_headers_provider(self) -> str: @@ -360,7 +360,7 @@ class AmazonAnthropicClaudeMessagesConfig( display: Final = thinking.get("display") if thinking is not None else None anthropic_messages_request["output_config"] = output_config anthropic_messages_request["thinking"] = { - **thinking_block_binding(thinking), + **(thinking_block_binding(thinking) if self._carries_thinking_block_binding else {}), "type": "adaptive", **({} if display is None else {"display": display}), } @@ -705,7 +705,7 @@ class AmazonAnthropicClaudeMessagesConfig( and translated_thinking.get("type") == "adaptive" ): anthropic_messages_request["thinking"] = { - **thinking_block_binding(translated_thinking), + **(thinking_block_binding(translated_thinking) if self._carries_thinking_block_binding else {}), "type": "adaptive", "display": "updates", } diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index 05aa3bfe5ef..01b9ca7d161 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -25,6 +25,9 @@ import pytest sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) +from litellm.llms.base_llm.anthropic_messages.transformation import ( # noqa: E402 # sys.path must be patched before importing litellm + BaseAnthropicMessagesConfig, +) from litellm.proxy._types import SpecialHeaders # noqa: E402 # sys.path must be patched before importing litellm # Fake tokens for testing (not real secrets) @@ -4193,10 +4196,21 @@ ADAPTIVE_THINKING_WITH_BLOCK_BINDING: Final = MappingProxyType( ("claude-opus-4-8", "azure_ai", ADAPTIVE_THINKING_WITH_BLOCK_BINDING), ("invoke/us.anthropic.claude-opus-4-8", "bedrock", ADAPTIVE_THINKING_WITH_BLOCK_BINDING), ("converse/us.anthropic.claude-opus-4-8", "bedrock", {"type": "adaptive"}), + ("mantle/anthropic.claude-opus-4-8", "bedrock", {"type": "adaptive"}), + ("claude_platform/claude-opus-4-8", "bedrock", {"type": "adaptive"}), ("claude-opus-4-8", "vertex_ai", {"type": "adaptive"}), ("databricks-claude-opus-4-8", "databricks", {"type": "adaptive"}), ), - ids=("anthropic", "azure_ai", "bedrock_invoke", "bedrock_converse", "vertex_ai", "databricks"), + ids=( + "anthropic", + "azure_ai", + "bedrock_invoke", + "bedrock_converse", + "bedrock_mantle", + "bedrock_claude_platform", + "vertex_ai", + "databricks", + ), ) def test_chat_legacy_thinking_translation_carries_only_block_binding_where_supported( model: str, custom_llm_provider: str, expected_thinking: Mapping[str, object] @@ -4219,14 +4233,24 @@ def test_chat_legacy_thinking_translation_carries_only_block_binding_where_suppo ("claude-opus-4-8", "anthropic", ADAPTIVE_THINKING_WITH_BLOCK_BINDING), ("claude-opus-4-8", "azure_ai", ADAPTIVE_THINKING_WITH_BLOCK_BINDING), ("us.anthropic.claude-opus-4-8", "bedrock", ADAPTIVE_THINKING_WITH_BLOCK_BINDING), + ("mantle/anthropic.claude-opus-4-8", "bedrock", {"type": "adaptive"}), + ("anthropic.claude-opus-4-8", "bedrock_mantle", {"type": "adaptive"}), + ("claude_platform/claude-opus-4-8", "bedrock", {"type": "adaptive"}), ("claude-opus-4-8", "vertex_ai", {"type": "adaptive"}), ), - ids=("anthropic", "azure_ai", "bedrock_invoke", "vertex_ai"), + ids=( + "anthropic", + "azure_ai", + "bedrock_invoke", + "bedrock_mantle", + "bedrock_mantle_provider", + "bedrock_claude_platform", + "vertex_ai", + ), ) def test_messages_legacy_thinking_translation_carries_only_block_binding_where_supported( model: str, custom_llm_provider: str, expected_thinking: Mapping[str, object] ) -> None: - from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager @@ -4235,7 +4259,22 @@ def test_messages_legacy_thinking_translation_carries_only_block_binding_where_s ) assert config is not None - result: Final = config.transform_anthropic_messages_request( + assert _thinking_sent_for_legacy_messages_request(config, model) == expected_thinking + + +@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") +def test_openai_compatible_messages_passthrough_keeps_legacy_thinking_without_block_binding() -> None: + from litellm.llms.openai_like.messages.transformation import OpenAILikeAnthropicMessagesConfig + + assert _thinking_sent_for_legacy_messages_request(OpenAILikeAnthropicMessagesConfig(), "claude-opus-4-8") == { + "type": "adaptive" + } + + +def _thinking_sent_for_legacy_messages_request(config: BaseAnthropicMessagesConfig, model: str) -> object: + from litellm.types.router import GenericLiteLLMParams + + return config.transform_anthropic_messages_request( model=model, messages=[{"role": "user", "content": "Reply with OK"}], anthropic_messages_optional_request_params={ @@ -4244,9 +4283,7 @@ def test_messages_legacy_thinking_translation_carries_only_block_binding_where_s }, litellm_params=GenericLiteLLMParams(), headers={}, - ) - - assert result.get("thinking") == expected_thinking + ).get("thinking") @pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index cc154c3d026..c572d48b773 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -747,7 +747,7 @@ def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_th "schema": {"type": "object", "properties": {"result": {"type": "string"}}}, }, }, - "thinking": {"type": "enabled", "budget_tokens": 4096}, + "thinking": {"type": "enabled", "budget_tokens": 4096, "block_binding": BLOCK_BINDING}, "max_tokens": 8192, }, optional_params={}, @@ -756,7 +756,7 @@ def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_th ) assert "tools" in result - assert result["thinking"] == {"type": "adaptive"} + assert result["thinking"] == {"type": "adaptive", "block_binding": BLOCK_BINDING} assert result["output_config"] == {"effort": "high"} @@ -1037,6 +1037,32 @@ def test_bedrock_chat_invoke_keeps_block_binding_with_display_updates( assert result.get("thinking") == {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING} +@pytest.mark.parametrize( + "config_path", + ( + "litellm.llms.bedrock.chat.mantle.transformation.AmazonMantleConfig", + "litellm.llms.bedrock_mantle.chat.claude_transformation.BedrockMantleClaudeChatConfig", + ), + ids=("bedrock_mantle_route", "bedrock_mantle_provider"), +) +def test_bedrock_chat_mantle_rebuilds_display_updates_without_block_binding( + local_model_cost_map, local_beta_headers_config, config_path: str +) -> None: + module_name, _, class_name = config_path.rpartition(".") + config: Final = getattr(importlib.import_module(module_name), class_name)() + optional_params: Final = config.map_openai_params( + non_default_params={ + "max_tokens": 512, + "thinking": {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, + }, + optional_params={}, + model="anthropic.claude-opus-4-8", + drop_params=False, + ) + + assert optional_params.get("thinking") == {"type": "adaptive", "display": "updates"} + + FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" EAGER_TOOL_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 059ef765990..11500db67bb 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -36,6 +36,8 @@ from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_tran AmazonAnthropicClaudeMessagesConfig, AmazonAnthropicClaudeMessagesStreamDecoder, ) +from litellm.llms.bedrock.messages.mantle_transformation import AmazonMantleMessagesConfig +from litellm.llms.bedrock_mantle.messages.transformation import BedrockMantleAnthropicMessagesConfig @pytest.mark.asyncio @@ -3536,15 +3538,31 @@ BLOCK_BINDING: Final = {"prefix_mismatch_behavior": "error"} @pytest.mark.parametrize( ("model", "requested_thinking", "extra_params"), ( - ("eu.anthropic.claude-opus-5", {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, {}), ( "eu.anthropic.claude-opus-5", - {"type": "enabled", "budget_tokens": 24000, "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, + {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, + {}, + ), + ( + "eu.anthropic.claude-opus-5", + { + "type": "enabled", + "budget_tokens": 24000, + "display": "updates", + "block_binding": BLOCK_BINDING, + "future_key": 1, + }, {}, ), ( "us.anthropic.claude-opus-4-6", - {"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, + { + "type": "enabled", + "budget_tokens": 2048, + "display": "updates", + "block_binding": BLOCK_BINDING, + "future_key": 1, + }, {"context_management": {"edits": [{"type": "clear_thinking_20251015"}]}}, ), ), @@ -3566,6 +3584,41 @@ def test_bedrock_messages_keeps_block_binding_with_display_updates( assert result.get("thinking") == {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING} +@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") +@pytest.mark.parametrize( + "config_class", + (AmazonMantleMessagesConfig, BedrockMantleAnthropicMessagesConfig), + ids=("bedrock_mantle_route", "bedrock_mantle_provider"), +) +@pytest.mark.parametrize( + ("requested_thinking", "extra_params"), + ( + ({"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, {}), + ( + {"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING}, + {"context_management": {"edits": [{"type": "clear_thinking_20251015"}]}}, + ), + ), + ids=("adaptive", "legacy_clear_thinking"), +) +def test_bedrock_mantle_messages_rebuilds_display_updates_without_block_binding( + config_class: type[AmazonMantleMessagesConfig], + requested_thinking: dict[str, object], + extra_params: dict[str, object], +) -> None: + from litellm.types.router import GenericLiteLLMParams + + result: Final = config_class().transform_anthropic_messages_request( + model="anthropic.claude-opus-4-8", + messages=[{"role": "user", "content": "Reply with OK"}], + anthropic_messages_optional_request_params={"max_tokens": 512, "thinking": requested_thinking, **extra_params}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("thinking") == {"type": "adaptive", "display": "updates"} + + @pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") @pytest.mark.parametrize("action", (None, "tool_addition", "tool_removal")) @pytest.mark.parametrize("explicit_beta", (False, True)) From 207148f3b6c3c87bb77c569db72606a3e1fb6a60 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 20 Oct 2026 02:42:05 -0700 Subject: [PATCH 13/13] style(openai): drop the unused Any import from the responses count handler The #47389 adaptation moved the handler's return type to dict[str, object] and kept the Any import, so ruff check reports F401 on this file at 1.106.1 and not at 1.106.0 --- litellm/llms/openai/responses/count_tokens/handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 6b34f04a36a..6f911e9c40f 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -6,7 +6,7 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint. import json from collections.abc import Mapping, Sequence -from typing import Any, Final +from typing import Final import httpx