diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index eb8e595c849..0967c22f93a 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -161,11 +161,58 @@ def _has_non_chat_user_block(message: Mapping[str, object]) -> bool: return any(isinstance(block, dict) and block.get("type") not in ValidUserMessageContentTypes for block in blocks) +_OPENAI_RESPONSES_BUILT_IN_TOOL_TYPES: Final = frozenset( + ( + "apply_patch", + "code_interpreter", + "computer", + "computer_use_preview", + "file_search", + "image_generation", + "local_shell", + "mcp", + "shell", + "web_search", + "web_search_2025_08_26", + "web_search_preview", + "web_search_preview_2025_03_11", + ) +) + + +_CHAT_INFERENCE_TRANSLATED_TOOL_TYPES: Final = frozenset(("namespace", "tool_search")) + + +def _is_responses_built_in_tool(tool: Mapping[str, object]) -> bool: + tool_type: Final = tool.get("type") + return isinstance(tool_type, str) and tool_type in _OPENAI_RESPONSES_BUILT_IN_TOOL_TYPES + + +def _is_chat_inference_translated_tool(tool: Mapping[str, object]) -> bool: + tool_type: Final = tool.get("type") + return isinstance(tool_type, str) and tool_type in _CHAT_INFERENCE_TRANSLATED_TOOL_TYPES + + +def _needs_tool_translation(tool: Mapping[str, object]) -> bool: + return _is_chat_tool(tool) or _is_chat_inference_translated_tool(tool) or _is_responses_built_in_tool(tool) + + def _gets_placeholder_name(tool: Mapping[str, object]) -> bool: from litellm.llms.anthropic.messages.adapters.transformation import is_provider_native_tool_dict name: Final = tool.get("name") - return not (_is_chat_tool(tool) or is_provider_native_tool_dict(tool) or (isinstance(name, str) and name.strip())) + return not ( + _is_chat_tool(tool) + or _is_chat_inference_translated_tool(tool) + or is_provider_native_tool_dict(tool) + or (isinstance(name, str) and name.strip()) + ) + + +def _chat_tools_for(tool: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for + + return (tool,) if _is_chat_inference_translated_tool(tool) else chat_completion_tools_for(tool) def _untranslatable_input( @@ -182,8 +229,8 @@ def _untranslatable_input( unnamed_tool: Final = next(filter(_gets_placeholder_name, tools), None) if unnamed_tool is not None: return ChatCountTranslationError( - message=f"Invalid request for token counting: tool type {unnamed_tool.get('type')!r} " - "has no Chat Completions equivalent" + message=f"Invalid request for token counting: LiteLLM cannot count a tool of type " + f"{unnamed_tool.get('type')!r} for this model" ) return None @@ -223,9 +270,7 @@ def _translated_chat_request( chat_params: Mapping[str, object], translator: _ChatRequestTranslator, ) -> _ChatParams: - from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for - - chat_tools: Final = [dict(chat_tool) for chat_tool in chain.from_iterable(map(chat_completion_tools_for, tools))] + chat_tools: Final = [dict(chat_tool) for chat_tool in chain.from_iterable(map(_chat_tools_for, tools))] openai_params: Final = {name: value for name, value in chat_params.items() if name != "drop_params"} mapped_params: Final = {**openai_params, **({"tools": chat_tools} if chat_tools else {})} chat_messages: Final = cast( # cast-ok: chat-shaped request messages, read the way chat inference reads them @@ -255,7 +300,7 @@ def anthropic_fields_for_chat_tool_use( counted_tools: Final = tools or () if ( not chat_params - and not any(map(_is_chat_tool, counted_tools)) + and not any(map(_needs_tool_translation, counted_tools)) and not any(map(_needs_chat_translation, messages)) ): return None diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 21fbdb04896..26bc5904e06 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -10,6 +10,7 @@ from typing import Any, Final from pydantic import JsonValue from typing_extensions import assert_never +import litellm from litellm._logging import verbose_logger from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.count_tokens.transformation import ChatCountTranslationError, anthropic_count_fields @@ -53,6 +54,18 @@ class CountTokensFailure: CountTokensOutcome = CountedTokens | CountTokensFailure +def _claude_count_translator(litellm_params: Mapping[str, object]) -> AnthropicConfig: + deployment_model: Final = litellm_params.get("model") + if not isinstance(deployment_model, str): + return AnthropicConfig() + provider_hint: Final = litellm_params.get("custom_llm_provider") + model, provider, _, _ = litellm.get_llm_provider( + model=deployment_model, custom_llm_provider=provider_hint if isinstance(provider_hint, str) else None + ) + chat_config: Final = litellm.ProviderConfigManager.get_provider_chat_config(model, LlmProviders(provider)) + return chat_config if isinstance(chat_config, AnthropicConfig) else AnthropicConfig() + + def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool: return ( isinstance(outcome, CountTokensFailure) @@ -138,16 +151,22 @@ class BedrockTokenCounter(BaseTokenCounter): counted_model: Final = model_to_use.removeprefix("bedrock_mantle/") resolved_model: Final = get_bedrock_base_model(counted_model) primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE + litellm_params: Final = (deployment or {}).get("litellm_params", {}) translated: Final = ( await anthropic_count_fields( - counted_model, messages, system, tools, messages_format, AnthropicConfig(), chat_params + counted_model, + messages, + system, + tools, + messages_format, + _claude_count_translator(litellm_params), + chat_params, ) if is_mantle_claude_model(resolved_model) else None ) if isinstance(translated, ChatCountTranslationError): return translated.response(request_model, model_to_use, primary_tokenizer) - litellm_params: Final = (deployment or {}).get("litellm_params", {}) runtime_request: Final = ( _request_data(model_to_use, messages, tools, system) if translated is None diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index 4e7dc012802..dd04e531f2d 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -73,6 +73,8 @@ class ProviderCountFailure: return not self.rejected_before_provider_call and self.status_code in REGION_DENIED_STATUSES def with_hint(self, hint: str) -> "ProviderCountFailure": + if self.rejected_before_provider_call: + return self return replace(self, message=_with_hint(self.message, hint)) def to_proxy_exception(self) -> ProxyException: diff --git a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py index 27cbbe81a13..380881a6796 100644 --- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py +++ b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py @@ -37,6 +37,7 @@ from integration.providers._count_tokens_system_lift import ( REMINDER, RESPONSES_REQUEST_OPTIONS, RESPONSES_TOOL_USE, + RESPONSES_TOOLS, STRING_CASE, TOKEN_COUNTING_BETA, TOOLS, @@ -603,7 +604,24 @@ def _generating(provider: _Provider) -> Callable[[Request], Reply]: return respond -@pytest.mark.parametrize("request_body", RESPONSES_TOOL_USE.values(), ids=RESPONSES_TOOL_USE.keys()) +_REWRITTEN_TOOL_USE: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + **RESPONSES_TOOL_USE, + "file-search": { + "input": WEATHER_QUESTION, + "tools": [{"type": "file_search", "vector_store_ids": ["vs_audit_weather"]}], + }, + "additional-tools": { + "input": [ + {"role": "user", "content": WEATHER_QUESTION}, + {"type": "additional_tools", "tools": RESPONSES_TOOLS}, + ], + }, + } +) + + +@pytest.mark.parametrize("request_body", _REWRITTEN_TOOL_USE.values(), ids=_REWRITTEN_TOOL_USE.keys()) def test_responses_input_tokens_counts_what_responses_inference_sends( deployment: _Deployment, gateway: Gateway, request_body: dict[str, JsonValue] ) -> None: @@ -642,6 +660,79 @@ def test_responses_input_tokens_counts_the_request_options_responses_inference_s assert count_request_fields(counted) == inference_count_fields(inference), (counted, inference) +def _without_cache_marks(value: JsonValue) -> JsonValue: + if isinstance(value, dict): + return {key: _without_cache_marks(item) for key, item in value.items() if key != "cache_control"} + if isinstance(value, list): + return [_without_cache_marks(item) for item in value] + return value + + +def test_responses_input_tokens_leaves_out_the_cache_marks_responses_inference_sends( + deployment: _Deployment, gateway: Gateway +) -> None: + marked: Final[dict[str, JsonValue]] = {"type": "ephemeral"} + request_body: Final[dict[str, JsonValue]] = { + "input": [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": f"Weather note {index}.", "cache_control": marked} + for index in range(2) + ], + }, + {"role": "user", "content": WEATHER_QUESTION}, + ], + "tools": RESPONSES_TOOLS, + } + with wire_server(_generating(deployment.provider), port=deployment.port) as peer: + generated: Final = gateway.request( + "POST", "/v1/responses", {**request_body, "model": deployment.model, "max_output_tokens": 16} + ) + assert generated.status_code == 200, generated.text + response: Final = gateway.request( + "POST", "/v1/responses/input_tokens", {**request_body, "model": deployment.model} + ) + assert _payload(response)["input_tokens"] == _COUNT, response.text + inference, counted = (counted_fields(body) for body in _inference_and_count(peer)) + assert "cache_control" in json.dumps(inference), inference + assert "cache_control" not in json.dumps(counted), counted + assert counted == _without_cache_marks(inference), (counted, inference) + + +_UNCOUNTABLE_TOOLS: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + "computer_use_preview": { + "type": "computer_use_preview", + "display_width": 1024, + "display_height": 768, + "environment": "browser", + }, + "mcp": { + "type": "mcp", + "server_label": "docs", + "server_url": "http://127.0.0.1:9/sse", + "require_approval": "never", + }, + } +) + + +@pytest.mark.parametrize("tool", _UNCOUNTABLE_TOOLS.values(), ids=_UNCOUNTABLE_TOOLS.keys()) +def test_responses_input_tokens_rejects_a_tool_it_cannot_count_without_calling_the_peer( + deployment: _Deployment, gateway: Gateway, tool: dict[str, JsonValue] +) -> None: + with _listening(deployment) as peer: + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": deployment.model, "input": WEATHER_QUESTION, "tools": [tool]}, + ) + assert response.status_code == 400, response.text + assert tool["type"] in response.text, response.text + assert peer.drain() == () + + def _text_blocks(message: JsonValue) -> JsonValue: if not isinstance(message, dict) or not isinstance(message.get("content"), str): return message diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index a4542f1f604..c392e355708 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -1,4 +1,5 @@ import asyncio +import base64 import json import re import threading @@ -32,6 +33,7 @@ from integration.providers._count_tokens_system_lift import ( LIFTED, MID_SYSTEM, REMINDER, + RESPONSES_REQUEST_OPTIONS, RESPONSES_TOOLS, STRING_CASE, USER, @@ -1702,6 +1704,72 @@ def test_responses_input_tokens_sends_mantle_the_tool_use_invoke_inference_sends assert _tool_use_fields(mantle_count) == expected, mantle_count +_RUNTIME_COUNTED_OPTIONS: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + {**RESPONSES_REQUEST_OPTIONS, "forced-tool-choice": _FORCED_TOOL_CHOICE} +) +_CONVERSE_WEATHER_TOOLS: Final[dict[str, JsonValue]] = { + "tools": [ + { + "toolSpec": { + "name": "get_weather", + "description": "Look up the current weather for a city", + "inputSchema": {"json": WEATHER_SCHEMA}, + } + } + ] +} +_INVOKE_WEATHER_TOOLS: Final[list[JsonValue]] = [ + { + "type": "custom", + "name": "get_weather", + "description": "Look up the current weather for a city", + "input_schema": WEATHER_SCHEMA, + } +] +_RUNTIME_COUNTED_SHAPES: Final[Mapping[str, tuple[frozenset[str], str, JsonValue]]] = MappingProxyType( + { + "converse": (frozenset({"messages", "system", "toolConfig"}), "toolConfig", _CONVERSE_WEATHER_TOOLS), + "invokeModel": ( + frozenset({"anthropic_version", "max_tokens", "messages", "system", "tools"}), + "tools", + _INVOKE_WEATHER_TOOLS, + ), + } +) + + +def _runtime_counted_request(count_body: Mapping[str, JsonValue]) -> tuple[str, dict[str, JsonValue]]: + count_input: Final = _JSON_OBJECT.validate_python(count_body["input"]) + invoke: Final = count_input.get("invokeModel") + if invoke is None: + return "converse", _JSON_OBJECT.validate_python(count_input["converse"]) + encoded: Final = _JSON_OBJECT.validate_python(invoke)["body"] + assert isinstance(encoded, str), invoke + return "invokeModel", _JSON_OBJECT.validate_json(base64.b64decode(encoded)) + + +@pytest.mark.parametrize("request_body", _RUNTIME_COUNTED_OPTIONS.values(), ids=_RUNTIME_COUNTED_OPTIONS.keys()) +def test_responses_input_tokens_sends_the_runtime_the_tools_without_the_options_it_refuses( + counting_proxy: OwnedProxy, mantle_port: int, request_body: dict[str, JsonValue] +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _metadata_deployment(scenario, runtime.url, _SONNET) + response: Final = gateway.request("POST", "/v1/responses/input_tokens", {**request_body, "model": model}) + assert _payload(response)["input_tokens"] == _RUNTIME_COUNT, response.text + ((target, count_body),) = _runtime_count_requests(runtime) + assert target == f"/model/{_SONNET}/count-tokens", target + shape, counted = _runtime_counted_request(count_body) + allowed_fields, tools_field, weather_tools = _RUNTIME_COUNTED_SHAPES[shape] + assert set(counted) <= allowed_fields, counted + assert counted.get(tools_field) == (weather_tools if "tools" in request_body else None), counted + assert mantle.drain() == () + + def _counting_every_model(request: Request) -> Reply: if _COUNT_TARGET.match(target_of(request)): return _json_reply(200, {"inputTokens": _RUNTIME_COUNT}) diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index 96410b03a92..9ecec3686a9 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -1,5 +1,7 @@ import json +import uuid from collections.abc import Iterator, Mapping +from dataclasses import dataclass from typing import Final import pytest @@ -488,3 +490,103 @@ def test_responses_input_tokens_sends_openai_the_owners_previous_response_id(cou assert counted.status_code == 200, counted.text count_body: Final = _body_for(wire.drain(), "/v1/responses/input_tokens") assert count_body["previous_response_id"] == _RESPONSES_REPLY["id"] + + +_PROVIDER_FILE_ID: Final = "file-audit-provider" +_UPLOADED_FILE: Final[Mapping[str, JsonValue]] = { + "id": _PROVIDER_FILE_ID, + "object": "file", + "bytes": 5, + "created_at": 1, + "filename": "notes.txt", + "purpose": "user_data", + "status": "processed", +} + + +def _files_peer(request: Request) -> Reply: + if request.target == "/v1/files": + return Reply(body=json.dumps(_UPLOADED_FILE).encode()) + return _count_response(request) + + +def _file_input(file_id: str) -> list[JsonValue]: + return [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize the notes."}, + {"type": "input_file", "file_id": file_id}, + ], + } + ] + + +@dataclass(frozen=True, slots=True) +class _FileTenants: + model_id: str + model: str + owner_key: str + outsider_key: str + file_id: str + + +def _file_tenants(scenario: Scenario, wire: Wire) -> _FileTenants: + model_id: Final = f"audit-openai-{uuid.uuid4().hex}" + model: Final = scenario.model( + model_info={"id": model_id}, + model="openai/gpt-4o-mini", + api_base=f"{wire.url}/v1", + api_key="sk-integration-openai", + ) + owner_key: Final = scenario.key( + user_id=scenario.user(user_role="internal_user"), team_id=scenario.team(models=[model]), models=[model] + ) + outsider_key: Final = scenario.key( + user_id=scenario.user(user_role="internal_user"), team_id=scenario.team(models=[model]), models=[model] + ) + uploaded: Final = scenario.gateway.request_multipart( + "/v1/files", + {"purpose": "user_data", "target_model_names": model}, + {"file": ("notes.txt", b"notes", "text/plain")}, + key=owner_key, + ) + assert uploaded.status_code == 200, uploaded.text + file_id: Final = _JSON_OBJECT.validate_json(uploaded.content)["id"] + assert isinstance(file_id, str) and file_id != _PROVIDER_FILE_ID, uploaded.text + assert [request.target for request in wire.drain()] == ["/v1/files"] + return _FileTenants(model_id, model, owner_key, outsider_key, file_id) + + +def test_responses_input_tokens_sends_openai_the_provider_id_of_a_managed_file(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_files_peer) as wire, gateway.scenario() as scenario: + tenants: Final = _file_tenants(scenario, wire) + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + { + "model": tenants.model, + "input": _file_input(tenants.file_id), + "model_file_id_mapping": {tenants.file_id: {tenants.model_id: "file-forged"}}, + }, + key=tenants.owner_key, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT + assert _count_body(wire) == {"model": "gpt-4o-mini", "input": _file_input(_PROVIDER_FILE_ID)} + + +def test_responses_input_tokens_refuses_another_teams_managed_file_before_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_files_peer) as wire, gateway.scenario() as scenario: + tenants: Final = _file_tenants(scenario, wire) + refused: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": tenants.model, "input": _file_input(tenants.file_id)}, + key=tenants.outsider_key, + ) + assert refused.status_code == 403, refused.text + assert _JSON_OBJECT.validate_json(refused.content)["error"]["type"] == "permission_error", refused.text + assert wire.drain() == () diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index 63f149fabe4..89352ce8702 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -569,6 +569,7 @@ _ANTHROPIC_TOOL_USE: Final[list[dict[str, JsonValue]]] = [ ("messages", "tool", "messages_format"), ( (_ANTHROPIC_TOOL_USE, {"name": "get_weather", "input_schema": _WEATHER_SCHEMA}, "openai"), + ([{"role": "user", "content": "Open example.com"}], {"type": "computer_toolset_20260801"}, "openai"), ( [{"role": "user", "content": "Weather in SF?"}], {"type": "function", "name": "get_weather", "parameters": _WEATHER_SCHEMA}, @@ -685,6 +686,26 @@ _CHAT_WEATHER_TOOL: Final[dict[str, JsonValue]] = { [{"type": "mcp", "server_label": "wiki", "server_url": "https://mcp.example/mcp"}, _CHAT_WEATHER_TOOL], id="mcp-tool", ), + pytest.param( + [{"role": "user", "content": "Open the docs"}], + [{"type": "computer_use_preview", "display_width": 1024, "display_height": 768, "environment": "browser"}], + id="computer-use-preview-tool-alone", + ), + pytest.param( + [{"role": "user", "content": "Ask the wiki"}], + [{"type": "mcp", "server_label": "wiki", "server_url": "https://mcp.example/mcp"}], + id="mcp-tool-alone", + ), + pytest.param( + [{"role": "user", "content": "Plot the sales data"}], + [{"type": "code_interpreter", "container": {"type": "auto"}}], + id="code-interpreter-tool-alone", + ), + pytest.param( + [{"role": "user", "content": "Fix the failing test"}], + [{"type": "apply_patch"}], + id="apply-patch-tool-alone", + ), ), ) async def test_counter_rejects_input_chat_translation_would_leave_out_without_calling_anthropic( @@ -700,6 +721,80 @@ async def test_counter_rejects_input_chat_translation_would_leave_out_without_ca assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +async def test_counter_sends_a_tool_with_a_non_string_type_to_anthropic_as_given( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], count_endpoint: str +) -> None: + tool: Final[dict[str, JsonValue]] = {"type": ["mcp"], "name": "lookup"} + + result, bodies = await _counted_body( + counter_type, count_endpoint, [{"role": "user", "content": "Ask the wiki"}], [tool], "openai" + ) + + assert result is not None + assert [body.get("tools") for body in bodies] == [[tool]] + + +_CRM_NAMESPACE_TOOL: Final[dict[str, JsonValue]] = { + "type": "namespace", + "name": "crm", + "description": "CRM lookups", + "tools": [ + { + "type": "function", + "name": "find_owner", + "description": "Find the owner of an account", + "parameters": {"type": "object", "properties": {"account": {"type": "string"}}, "required": ["account"]}, + } + ], +} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +@pytest.mark.parametrize( + "tools", + ( + pytest.param([{"type": "tool_search"}], id="tool-search-alone"), + pytest.param( + [ + {"type": "tool_search"}, + { + "type": "function", + "function": {"name": "get_weather", "parameters": _WEATHER_SCHEMA}, + "defer_loading": True, + }, + ], + id="tool-search-with-deferred-function", + ), + pytest.param([_CRM_NAMESPACE_TOOL], id="namespace-alone"), + pytest.param([_CHAT_WEATHER_TOOL, _CRM_NAMESPACE_TOOL], id="namespace-with-function"), + ), +) +async def test_counter_sends_the_tools_chat_inference_makes_from_tool_search_and_namespace( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], + count_endpoint: str, + tools: list[dict[str, JsonValue]], +) -> None: + messages: Final[list[dict[str, JsonValue]]] = [{"role": "user", "content": "Who owns the north region account?"}] + config: Final = AnthropicConfig() + inference: Final = config.transform_request( + model="claude-opus-5-5", + messages=deepcopy(messages), + optional_params=config.map_openai_params({"tools": deepcopy(tools)}, {}, "claude-opus-5-5", drop_params=False), + litellm_params={}, + headers={}, + ) + + result, bodies = await _counted_body(counter_type, count_endpoint, messages, tools, "openai") + + assert result is not None + assert [body.get("tools") for body in bodies] == [inference["tools"]] + + _ANSWER_SCHEMA: Final[dict[str, JsonValue]] = { "type": "object", "properties": {"answer": {"type": "string"}}, diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index 6fb47491323..cec7950b859 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -6,6 +6,7 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter +import litellm from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -264,6 +265,79 @@ async def test_claude_chat_tool_use_is_counted_in_the_shape_claude_inference_sen {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "sunny"}]}, ] + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool", + ( + {"type": "mcp", "server_label": "wiki", "server_url": "https://mcp.example/mcp"}, + {"type": "computer_use_preview", "display_width": 1024, "display_height": 768, "environment": "browser"}, + ), + ids=("mcp", "computer_use_preview"), +) +async def test_claude_count_rejects_a_lone_responses_tool_without_calling_bedrock(tool: dict[str, JsonValue]) -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=[{"role": "user", "content": "Ask the wiki"}], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + tools=[tool], + messages_format="openai", + ) + + assert result is not None + assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) + assert bedrock.posted_hosts() == () + + +_CLAUDE_REPLY: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 3, "output_tokens": 1}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment_model", + ("bedrock/invoke/us.anthropic.claude-sonnet-5-5", "bedrock_mantle/us.anthropic.claude-sonnet-5-5"), +) +async def test_claude_count_sends_the_tool_search_tool_the_deployments_inference_sends(deployment_model: str) -> None: + credentials: Final = DEPLOYMENT["litellm_params"] + inference: Final = _Bedrock(runtime=(200, _CLAUDE_REPLY), mantle=(200, _CLAUDE_REPLY)) + await litellm.acompletion( + model=deployment_model, + messages=MESSAGES, + tools=[{"type": "tool_search"}], + max_tokens=8, + client=AsyncHTTPHandler(transport=httpx.MockTransport(inference)), + **credentials, + ) + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="bedrock_mantle/us.anthropic.claude-opus-5-5", + messages=MESSAGES, + contents=None, + deployment={"litellm_params": {**credentials, "model": deployment_model}}, + request_model="claude-sonnet-5-5", + tools=[{"type": "tool_search"}], + messages_format="openai", + ) + + assert result is not None + assert (result.error, result.total_tokens) == (False, 2177) + counted_tools: Final = _JSON_BODY.validate_json(bedrock.requests[0].content)["tools"] + assert counted_tools == _JSON_BODY.validate_json(inference.requests[0].content)["tools"] + + @pytest.mark.asyncio async def test_claude_chat_tool_use_counted_on_mantle_metadata_route_is_translated_too() -> None: bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py index 7ff32f9d4f3..43edc7e449f 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -229,3 +229,10 @@ def test_hint_lands_in_the_upstream_error_message_and_keeps_its_envelope(message assert (json.loads(hinted.message) if isinstance(expected, dict) else hinted.message) == expected assert hinted.status_code == 404 + + +def test_hint_stays_off_an_error_raised_before_any_provider_call() -> None: + message: Final = "Invalid request for token counting: LiteLLM cannot count a tool of type 'mcp' for this model" + failure: Final = ProviderCountFailure(status_code=400, message=message, rejected_before_provider_call=True) + + assert failure.with_hint("Enable eu-west-1").message == message