mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
fix(anthropic): keep tool names over 64 chars when a guardrail rewrites a /v1/messages request (#47434)
* fix(anthropic): keep tool names over 64 chars when a guardrail rewrites a /v1/messages request The guardrail translation shortens long tool names for the OpenAI-shaped guardrail input and wrote the shortened names back into tools and tool_use blocks, while tool_reference blocks and tool_choice kept the original names, so Bedrock rejected the request * fix(anthropic): restore declared tool names only from the request's own tool list A history tool_use name whose shortened form matched a declared short tool renamed that declared tool, so a forced tool_choice no longer matched and a tool allowlist saw a different name than the one forwarded * fix(anthropic): keep a tool_use name the caller already used as written A history tool_use named exactly like the shortened form of another long name was rewritten to that long name; such an ambiguous alias is now left alone --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
This commit is contained in:
yucheng-berri
GitHub
parent
6564a490f4
commit
78a6eeb212
4 changed files with 1105 additions and 2 deletions
|
|
@ -13,7 +13,7 @@ Pattern Overview:
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, MutableSequence, Sequence
|
||||
from collections.abc import Iterable, Mapping, MutableSequence, Sequence
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain, repeat
|
||||
|
|
@ -27,8 +27,10 @@ from litellm._logging import verbose_proxy_logger
|
|||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.pass_through.adapters.transformation import (
|
||||
OPENAI_MAX_TOOL_NAME_LENGTH,
|
||||
LiteLLMAnthropicMessagesAdapter,
|
||||
is_provider_native_tool_dict,
|
||||
truncate_tool_name,
|
||||
)
|
||||
from litellm.llms.anthropic.pass_through.convertible_blocks import (
|
||||
keep_non_string_tool_result_texts,
|
||||
|
|
@ -202,6 +204,60 @@ def _is_client_tool_use(block: Mapping[str, object]) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def _tool_use_name(block: object) -> str | None:
|
||||
fields: Final = as_str_mapping(block)
|
||||
if fields is None or fields.get("type") != "tool_use":
|
||||
return None
|
||||
return anthropic_tool_name(fields)
|
||||
|
||||
|
||||
def _tool_use_names(message: object) -> tuple[str, ...]:
|
||||
return tuple(name for name in map(_tool_use_name, list_field(message, "content")) if name is not None)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _OriginalToolNames:
|
||||
"""Caller tool names over the guardrail's limit, keyed by the shortened name the guardrail sees."""
|
||||
|
||||
declared: Mapping[str, str]
|
||||
used: Mapping[str, str]
|
||||
|
||||
|
||||
def _restorable_names(names: Iterable[str]) -> Mapping[str, str]:
|
||||
listed: Final = tuple(names)
|
||||
taken: Final = frozenset(listed)
|
||||
shortened: Final = tuple(
|
||||
(truncate_tool_name(name), name) for name in listed if len(name) > OPENAI_MAX_TOOL_NAME_LENGTH
|
||||
)
|
||||
return MappingProxyType({alias: name for alias, name in shortened if alias not in taken})
|
||||
|
||||
|
||||
def _original_tool_names(request: Mapping[str, object]) -> _OriginalToolNames:
|
||||
return _OriginalToolNames(
|
||||
declared=_restorable_names(chain.from_iterable(map(anthropic_tool_names, list_field(request, "tools")))),
|
||||
used=_restorable_names(chain.from_iterable(map(_tool_use_names, list_field(request, "messages")))),
|
||||
)
|
||||
|
||||
|
||||
def _restore_tool_use_names(messages: Sequence[object], original_names: Mapping[str, str]) -> None:
|
||||
for message in messages:
|
||||
for block in list_field(message, "content"):
|
||||
if (name := _tool_use_name(block)) is not None and name in original_names and isinstance(block, dict):
|
||||
block["name"] = original_names[name]
|
||||
|
||||
|
||||
def _with_original_name(tool: ChatCompletionToolParam, original_names: Mapping[str, str]) -> ChatCompletionToolParam:
|
||||
if "function" not in tool:
|
||||
return tool
|
||||
function: Final = tool["function"]
|
||||
fields: Final = as_str_mapping(function)
|
||||
name: Final = None if fields is None else fields.get("name")
|
||||
original: Final = original_names.get(name) if isinstance(name, str) else None
|
||||
if original is None:
|
||||
return tool
|
||||
return {**tool, "function": {**function, "name": original}}
|
||||
|
||||
|
||||
def _write_back_system_block(system: object, block_idx: int, response: str) -> None:
|
||||
if not isinstance(system, list):
|
||||
return
|
||||
|
|
@ -629,6 +685,9 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
model: Final = data.get("model")
|
||||
if model:
|
||||
inputs["model"] = model
|
||||
original_tool_names: Final = _original_tool_names(
|
||||
cast(Mapping[str, object], data) # cast-ok: Anthropic request body
|
||||
)
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=data,
|
||||
|
|
@ -643,7 +702,9 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
anthropic_config: Final = AnthropicConfig()
|
||||
anthropic_tools: Final[list[AllAnthropicToolsValues]] = []
|
||||
for tool in guardrailed_tools:
|
||||
converted_tool, mcp_server = anthropic_config.map_tool_helper(tool)
|
||||
converted_tool, mcp_server = anthropic_config.map_tool_helper(
|
||||
_with_original_name(tool, original_tool_names.declared)
|
||||
)
|
||||
if converted_tool is not None:
|
||||
anthropic_tools.append(converted_tool)
|
||||
# Note: MCP servers are handled separately in the main transformation
|
||||
|
|
@ -676,6 +737,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
),
|
||||
hoisted_system_message=hoisted_system_message,
|
||||
preserve_system_messages=has_midturn_system_message,
|
||||
original_tool_names=original_tool_names.used,
|
||||
)
|
||||
else:
|
||||
if guardrailed_texts and len(guardrailed_texts) != len(scanned):
|
||||
|
|
@ -840,6 +902,8 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
structured_messages: list, # mutable-ok: API message payload
|
||||
hoisted_system_message: object = None,
|
||||
preserve_system_messages: bool = False,
|
||||
*,
|
||||
original_tool_names: Mapping[str, str],
|
||||
) -> None:
|
||||
"""Write a guardrail's structured-message rewrite back without losing corrections."""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
|
|
@ -903,6 +967,10 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
for block in content:
|
||||
if isinstance(block, dict) and block.get("type") == "thinking":
|
||||
block.pop("cache_control", None)
|
||||
_restore_tool_use_names(
|
||||
cast(Sequence[object], converted), # cast-ok: Anthropic message payload
|
||||
original_tool_names,
|
||||
)
|
||||
data["messages"] = keep_non_string_tool_result_texts(converted, data.get("messages"))
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -0,0 +1,312 @@
|
|||
import base64
|
||||
import json
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, gateway_from_environment, object_value
|
||||
from integration._support.ports import allocate_port
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
from integration._support.upstream import _aws_event_frame
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(4 * graceful_stop_seconds() + 240)
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_TOKEN: Final = "synthetic-bedrock-bearer"
|
||||
_GUARDRAIL: Final = "headroom-compression"
|
||||
_SONNET: Final = "claude-sonnet-5-5"
|
||||
_OPUS: Final = "claude-opus-5"
|
||||
_SEARCH: Final = "ToolSearch"
|
||||
_NAME_SOURCE: Final = "mcp__plugin_support-kit_remote-workspace__getAccessibleWorkspaceResourcesForTheCurrentUser"
|
||||
_REPLY_TEXT: Final = "done"
|
||||
_FORCED_ID: Final = "toolu_forced"
|
||||
|
||||
|
||||
def _tool_name(length: int) -> str:
|
||||
name: Final = _NAME_SOURCE[:length]
|
||||
assert len(name) == length, name
|
||||
return name
|
||||
|
||||
|
||||
def _tool(name: str, **fields: JsonValue) -> dict[str, JsonValue]:
|
||||
return {"name": name, "description": "Returns OK", "input_schema": {"type": "object", "properties": {}}, **fields}
|
||||
|
||||
|
||||
def _blocks(message: JsonValue) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
content: Final = object_value(message).get("content")
|
||||
return tuple(object_value(block) for block in content) if isinstance(content, list) else ()
|
||||
|
||||
|
||||
def _parts(block: Mapping[str, JsonValue]) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
content: Final = block.get("content")
|
||||
return tuple(object_value(part) for part in content) if isinstance(content, list) else ()
|
||||
|
||||
|
||||
def _tool_results(blocks: tuple[Mapping[str, JsonValue], ...]) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
return tuple(block for block in blocks if block.get("type") == "tool_result")
|
||||
|
||||
|
||||
def _message_blocks(body: Mapping[str, JsonValue]) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
messages: Final = body.get("messages")
|
||||
assert isinstance(messages, list), body
|
||||
return tuple(block for message in messages for block in _blocks(message))
|
||||
|
||||
|
||||
def _tool_references(body: Mapping[str, JsonValue]) -> tuple[str, ...]:
|
||||
parts: Final = tuple(part for block in _tool_results(_message_blocks(body)) for part in _parts(block))
|
||||
return tuple(str(part["tool_name"]) for part in parts if part.get("type") == "tool_reference")
|
||||
|
||||
|
||||
def _referenced_names(body: Mapping[str, JsonValue]) -> tuple[str, ...]:
|
||||
used: Final = tuple(str(block["name"]) for block in _message_blocks(body) if block.get("type") == "tool_use")
|
||||
return (*used, *_tool_references(body))
|
||||
|
||||
|
||||
def _tool_names(body: Mapping[str, JsonValue]) -> tuple[str, ...]:
|
||||
tools: Final = body.get("tools")
|
||||
assert isinstance(tools, list), body
|
||||
return tuple(str(object_value(tool)["name"]) for tool in tools)
|
||||
|
||||
|
||||
def _forced_name(body: Mapping[str, JsonValue]) -> str | None:
|
||||
choice: Final = body.get("tool_choice")
|
||||
if not isinstance(choice, dict) or choice.get("type") != "tool":
|
||||
return None
|
||||
return str(choice["name"])
|
||||
|
||||
|
||||
def _bedrock_error(message: str) -> Reply:
|
||||
return Reply(status=400, body=json.dumps({"message": message}).encode())
|
||||
|
||||
|
||||
def _stop_reason(forced: str | None) -> str:
|
||||
return "end_turn" if forced is None else "tool_use"
|
||||
|
||||
|
||||
def _content(forced: str | None) -> list[JsonValue]:
|
||||
if forced is None:
|
||||
return [{"type": "text", "text": _REPLY_TEXT}]
|
||||
return [{"type": "tool_use", "id": _FORCED_ID, "name": forced, "input": {}}]
|
||||
|
||||
|
||||
def _message(model: str, forced: str | None) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": "msg_long_tool_names",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": model,
|
||||
"content": _content(forced),
|
||||
"stop_reason": _stop_reason(forced),
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 11, "output_tokens": 3},
|
||||
}
|
||||
|
||||
|
||||
def _invoke_chunk(event: Mapping[str, JsonValue]) -> bytes:
|
||||
payload: Final = {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}
|
||||
return _aws_event_frame("chunk", payload, "sc", "u")
|
||||
|
||||
|
||||
def _block_events(forced: str | None) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
if forced is None:
|
||||
return (
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _REPLY_TEXT}},
|
||||
)
|
||||
tool_use: Final = {"type": "tool_use", "id": _FORCED_ID, "name": forced, "input": {}}
|
||||
return (
|
||||
{"type": "content_block_start", "index": 0, "content_block": tool_use},
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{}"}},
|
||||
)
|
||||
|
||||
|
||||
def _stream(model: str, forced: str | None) -> Reply:
|
||||
start: Final = {**_message(model, forced), "content": [], "stop_reason": None}
|
||||
events: Final = (
|
||||
{"type": "message_start", "message": start},
|
||||
*_block_events(forced),
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{"type": "message_delta", "delta": {"stop_reason": _stop_reason(forced)}, "usage": {"output_tokens": 3}},
|
||||
{"type": "message_stop"},
|
||||
)
|
||||
frames: Final = tuple(_invoke_chunk(event) for event in events)
|
||||
return Reply(chunks=frames, content_type="application/vnd.amazon.eventstream")
|
||||
|
||||
|
||||
def _bedrock(request: Request) -> Reply:
|
||||
body: Final = _JSON_OBJECT.validate_json(request.body)
|
||||
available: Final = frozenset(_tool_names(body))
|
||||
missing: Final = [name for name in _tool_references(body) if name not in available]
|
||||
if missing:
|
||||
return _bedrock_error(f"Tool reference '{missing[0]}' not found in available tools")
|
||||
forced: Final = _forced_name(body)
|
||||
if forced is not None and forced not in available:
|
||||
return _bedrock_error(f"Tool '{forced}' not found in provided tools")
|
||||
model: Final = request.target.split("/")[2]
|
||||
if request.target.endswith("/invoke-with-response-stream"):
|
||||
return _stream(model, forced)
|
||||
return Reply(body=json.dumps(_message(model, forced)).encode())
|
||||
|
||||
|
||||
def _compress(request: Request) -> Reply:
|
||||
assert request.target == "/v1/compress", request.target
|
||||
messages: Final = _JSON_OBJECT.validate_json(request.body)["messages"]
|
||||
reply: Final = {"messages": messages, "tokens_before": 40, "tokens_after": 32, "compression_ratio": 0.8}
|
||||
return Reply(body=json.dumps(reply).encode())
|
||||
|
||||
|
||||
def _deployments(provider_url: str) -> list[JsonValue]:
|
||||
shared: Final = {"aws_region_name": "eu-central-1", "api_base": provider_url, "api_key": _TOKEN, "guardrails": []}
|
||||
return [
|
||||
{
|
||||
"model_name": _SONNET,
|
||||
"litellm_params": {
|
||||
"model": "global.anthropic.claude-sonnet-5-5",
|
||||
"custom_llm_provider": "bedrock",
|
||||
**shared,
|
||||
},
|
||||
},
|
||||
{"model_name": _OPUS, "litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5", **shared}},
|
||||
]
|
||||
|
||||
|
||||
def _config(directory: Path, provider_url: str, headroom_url: str) -> Path:
|
||||
stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
guardrail: Final = {
|
||||
"guardrail": "headroom",
|
||||
"mode": "pre_call",
|
||||
"default_on": True,
|
||||
"api_base": headroom_url,
|
||||
"unreachable_fallback": "fail_open",
|
||||
}
|
||||
config: Final[dict[str, JsonValue]] = {
|
||||
**stock,
|
||||
"model_list": _deployments(provider_url),
|
||||
"router_settings": {**object_value(stock["router_settings"]), "num_retries": 0},
|
||||
"guardrails": [{"guardrail_name": _GUARDRAIL, "litellm_params": guardrail}],
|
||||
}
|
||||
path: Final = directory / "long-tool-names.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
gateway: Gateway
|
||||
provider: Wire
|
||||
compressor: Wire | None
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=("fail_open", "compressed"))
|
||||
def rig(request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp(f"long-tool-names-{request.param}")
|
||||
with wire_server(_bedrock) as provider, wire_server(_compress) as compressor, gateway_from_environment() as gateway:
|
||||
reachable: Final = request.param == "compressed"
|
||||
headroom_url: Final = compressor.url if reachable else f"http://127.0.0.1:{allocate_port()}"
|
||||
config: Final = _config(directory, provider.url, headroom_url)
|
||||
with owned_proxy(gateway, directory, {}, config=config, workers=2) as owned:
|
||||
yield _Rig(owned, provider, compressor if reachable else None)
|
||||
|
||||
|
||||
def _search_turns(name: str) -> list[JsonValue]:
|
||||
return [
|
||||
{"role": "user", "content": "find the tool"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "tool_use", "id": "toolu_01", "name": _SEARCH, "input": {"query": "x"}}],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_01",
|
||||
"content": [{"type": "tool_reference", "tool_name": name}],
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": "Found it."},
|
||||
{"role": "user", "content": "say done"},
|
||||
]
|
||||
|
||||
|
||||
def _history_turns(name: str) -> list[JsonValue]:
|
||||
return [
|
||||
{"role": "user", "content": "use the tool"},
|
||||
{"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_02", "name": name, "input": {}}]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "tool_result", "tool_use_id": "toolu_02", "content": "OK"},
|
||||
{"type": "text", "text": "say done"},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _forced_turns() -> list[JsonValue]:
|
||||
return [
|
||||
{"role": "user", "content": "hello"},
|
||||
{"role": "assistant", "content": "hi"},
|
||||
{"role": "user", "content": "call the tool"},
|
||||
]
|
||||
|
||||
|
||||
def _body(shape: str, model: str, name: str, stream: bool) -> dict[str, JsonValue]:
|
||||
tools: Final = [_tool(_SEARCH), _tool(name, defer_loading=True)] if shape == "tool_reference" else [_tool(name)]
|
||||
body: Final[dict[str, JsonValue]] = {"model": model, "max_tokens": 20, "stream": stream, "tools": tools}
|
||||
match shape:
|
||||
case "tool_reference":
|
||||
return {**body, "messages": _search_turns(name)}
|
||||
case "tool_use_history":
|
||||
return {**body, "messages": _history_turns(name)}
|
||||
case _:
|
||||
return {**body, "messages": _forced_turns(), "tool_choice": {"type": "tool", "name": name}}
|
||||
|
||||
|
||||
def _sse_tool_names(text: str) -> tuple[str, ...]:
|
||||
events: Final = tuple(
|
||||
_JSON_OBJECT.validate_json(line.removeprefix("data: "))
|
||||
for line in text.splitlines()
|
||||
if line.startswith("data: ")
|
||||
)
|
||||
starts: Final = tuple(
|
||||
object_value(event["content_block"]) for event in events if event.get("type") == "content_block_start"
|
||||
)
|
||||
return tuple(str(block["name"]) for block in starts if block.get("type") == "tool_use")
|
||||
|
||||
|
||||
def _answered_tool_names(text: str, stream: bool) -> tuple[str, ...]:
|
||||
if stream:
|
||||
return _sse_tool_names(text)
|
||||
content: Final = _JSON_OBJECT.validate_json(text)["content"]
|
||||
assert isinstance(content, list), text
|
||||
return tuple(str(block["name"]) for block in map(object_value, content) if block.get("type") == "tool_use")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", (False, True), ids=("json", "sse"))
|
||||
@pytest.mark.parametrize("length", (64, 65, 90))
|
||||
@pytest.mark.parametrize("shape", ("tool_reference", "tool_choice", "tool_use_history"))
|
||||
def test_guarded_messages_keep_long_tool_names_the_caller_sent(
|
||||
rig: _Rig, shape: str, length: int, stream: bool
|
||||
) -> None:
|
||||
name: Final = _tool_name(length)
|
||||
body: Final = _body(shape, _OPUS if shape == "tool_choice" else _SONNET, name, stream)
|
||||
response: Final = rig.gateway.request("POST", "/v1/messages", body, headers={"anthropic-version": "2023-06-01"})
|
||||
sent: Final = rig.provider.drain()
|
||||
compressed: Final = None if rig.compressor is None else rig.compressor.drain()
|
||||
|
||||
assert _GUARDRAIL in response.headers.get("x-litellm-applied-guardrails", ""), dict(response.headers)
|
||||
assert compressed is None or len(compressed) == 1, compressed
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(sent) == 1, sent
|
||||
outbound: Final = _JSON_OBJECT.validate_json(sent[0].body)
|
||||
assert _tool_names(outbound) == ((_SEARCH, name) if shape == "tool_reference" else (name,)), outbound
|
||||
assert _referenced_names(outbound) == _referenced_names(body), outbound
|
||||
expected_calls: Final = (name,) if shape == "tool_choice" else ()
|
||||
assert _answered_tool_names(response.text, stream) == expected_calls, response.text
|
||||
|
|
@ -0,0 +1,585 @@
|
|||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.anthropic_sse import sse
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
|
||||
from integration._support.ports import allocate_port
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from litellm.llms.anthropic.pass_through.adapters.transformation import truncate_tool_name
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_JSON_OBJECTS: Final = TypeAdapter(list[dict[str, JsonValue]])
|
||||
_NAME_SOURCE: Final = (
|
||||
"mcp__plugin_support-kit_remote-workspace__getAccessibleWorkspaceResourcesForTheCurrentUser"
|
||||
"_including_archived_and_shared_workspaces_across_every_region_and_every_tenant_the_user_belongs_to"
|
||||
"_sorted_by_last_activity_with_owner_and_role_details"
|
||||
)
|
||||
_SEARCH: Final = "ToolSearch"
|
||||
_SHORT_NAME: Final = "workspace_resources"
|
||||
_EMAIL: Final = "jane.doe@example.com"
|
||||
_MASK: Final = "<EMAIL>"
|
||||
_FORCED_ID: Final = "toolu_forced"
|
||||
_GUARDRAIL_NAME_LIMIT: Final = 64
|
||||
_BLOCK_REASON: Final = "tool policy says no"
|
||||
|
||||
|
||||
def _tool_name(length: int) -> str:
|
||||
name: Final = _NAME_SOURCE[:length]
|
||||
assert len(name) == length, name
|
||||
return name
|
||||
|
||||
|
||||
def _tool(name: str) -> dict[str, JsonValue]:
|
||||
return {"name": name, "description": "Returns OK", "input_schema": {"type": "object", "properties": {}}}
|
||||
|
||||
|
||||
def _every_carrier(model: str, name: str, guardrail: str, stream: bool) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model": model,
|
||||
"max_tokens": 32,
|
||||
"stream": stream,
|
||||
"guardrails": [guardrail],
|
||||
"tools": [_tool(_SEARCH), {**_tool(name), "defer_loading": True}],
|
||||
"tool_choice": {"type": "tool", "name": name},
|
||||
"messages": [
|
||||
{"role": "user", "content": "find the workspace tool"},
|
||||
{"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": "tool_reference", "tool_name": name}],
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_02", "name": name, "input": {}}]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "tool_result", "tool_use_id": "toolu_02", "content": "workspace-alpha"},
|
||||
{"type": "text", "text": f"share it with {_EMAIL}"},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _blocks(body: Mapping[str, JsonValue]) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
messages: Final = _JSON_OBJECTS.validate_python(body["messages"])
|
||||
contents: Final = tuple(message["content"] for message in messages if isinstance(message["content"], list))
|
||||
return tuple(object_value(block) for content in contents for block in content)
|
||||
|
||||
|
||||
def _tool_result_parts(body: Mapping[str, JsonValue]) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
contents: Final = tuple(block.get("content") for block in _blocks(body) if block.get("type") == "tool_result")
|
||||
return tuple(object_value(part) for content in contents if isinstance(content, list) for part in content)
|
||||
|
||||
|
||||
def _named_tools(body: Mapping[str, JsonValue]) -> tuple[str, ...]:
|
||||
return tuple(string_value(tool["name"]) for tool in _JSON_OBJECTS.validate_python(body.get("tools") or []))
|
||||
|
||||
|
||||
def _tool_use_names(body: Mapping[str, JsonValue]) -> tuple[str, ...]:
|
||||
return tuple(string_value(block["name"]) for block in _blocks(body) if block.get("type") == "tool_use")
|
||||
|
||||
|
||||
def _tool_references(body: Mapping[str, JsonValue]) -> tuple[str, ...]:
|
||||
parts: Final = _tool_result_parts(body)
|
||||
return tuple(string_value(part["tool_name"]) for part in parts if part.get("type") == "tool_reference")
|
||||
|
||||
|
||||
def _forced_name(body: Mapping[str, JsonValue]) -> str | None:
|
||||
choice: Final = body.get("tool_choice")
|
||||
if not isinstance(choice, dict) or choice.get("type") != "tool":
|
||||
return None
|
||||
return string_value(choice["name"])
|
||||
|
||||
|
||||
def _carriers(body: Mapping[str, JsonValue]) -> tuple[tuple[str, ...], str | None, tuple[str, ...], tuple[str, ...]]:
|
||||
return _named_tools(body), _forced_name(body), _tool_use_names(body), _tool_references(body)
|
||||
|
||||
|
||||
def _anthropic_error(message: str) -> Reply:
|
||||
error: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": message}}
|
||||
return Reply(status=400, body=json.dumps(error).encode())
|
||||
|
||||
|
||||
def _message(model: str, forced: str | None) -> dict[str, JsonValue]:
|
||||
content: Final[list[JsonValue]] = (
|
||||
[{"type": "text", "text": "done"}]
|
||||
if forced is None
|
||||
else [{"type": "tool_use", "id": _FORCED_ID, "name": forced, "input": {}}]
|
||||
)
|
||||
return {
|
||||
"id": f"msg_{uuid.uuid4().hex}",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": model,
|
||||
"content": content,
|
||||
"stop_reason": "end_turn" if forced is None else "tool_use",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 11, "output_tokens": 3},
|
||||
}
|
||||
|
||||
|
||||
def _sse_reply(message: Mapping[str, JsonValue]) -> Reply:
|
||||
block: Final = object_value(_JSON_OBJECTS.validate_python(message["content"])[0])
|
||||
delta: Final[dict[str, JsonValue]] = (
|
||||
{"type": "input_json_delta", "partial_json": "{}"}
|
||||
if block["type"] == "tool_use"
|
||||
else {"type": "text_delta", "text": string_value(block["text"])}
|
||||
)
|
||||
started: Final = {**block, "input": {}} if block["type"] == "tool_use" else {**block, "text": ""}
|
||||
events: Final = (
|
||||
sse("message_start", {"type": "message_start", "message": {**message, "content": [], "stop_reason": None}}),
|
||||
sse("content_block_start", {"type": "content_block_start", "index": 0, "content_block": started}),
|
||||
sse("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": delta}),
|
||||
sse("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
sse(
|
||||
"message_delta",
|
||||
{"type": "message_delta", "delta": {"stop_reason": message["stop_reason"]}, "usage": {"output_tokens": 3}},
|
||||
),
|
||||
sse("message_stop", {"type": "message_stop"}),
|
||||
)
|
||||
return Reply(chunks=events, content_type="text/event-stream")
|
||||
|
||||
|
||||
def _anthropic_peer(request: Request) -> Reply:
|
||||
assert request.target.endswith("/v1/messages"), request.target
|
||||
body: Final = _JSON_OBJECT.validate_json(request.body)
|
||||
available: Final = frozenset(_named_tools(body))
|
||||
missing: Final = [name for name in _tool_references(body) if name not in available]
|
||||
if missing:
|
||||
return _anthropic_error(f"Tool reference '{missing[0]}' not found in available tools")
|
||||
forced: Final = _forced_name(body)
|
||||
if forced is not None and forced not in available:
|
||||
return _anthropic_error(f"Tool '{forced}' not found in provided tools")
|
||||
message: Final = _message(string_value(body["model"]), forced)
|
||||
return _sse_reply(message) if body.get("stream") is True else Reply(body=json.dumps(message).encode())
|
||||
|
||||
|
||||
def _chat_peer(request: Request) -> Reply:
|
||||
assert request.target.endswith("/chat/completions"), request.target
|
||||
body: Final = _JSON_OBJECT.validate_json(request.body)
|
||||
function: Final = object_value(_JSON_OBJECTS.validate_python(body["tools"])[0]["function"])
|
||||
call: Final = {"id": "call_1", "type": "function", "function": {"name": function["name"], "arguments": "{}"}}
|
||||
reply: Final = {
|
||||
"id": "chatcmpl-long-tool-names",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": body["model"],
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": None, "tool_calls": [call]},
|
||||
"finish_reason": "tool_calls",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 11, "completion_tokens": 3, "total_tokens": 14},
|
||||
}
|
||||
return Reply(body=json.dumps(reply).encode())
|
||||
|
||||
|
||||
def _masked(content: JsonValue) -> JsonValue:
|
||||
if isinstance(content, str):
|
||||
return content.replace(_EMAIL, _MASK)
|
||||
if isinstance(content, list):
|
||||
return [_masked(part) for part in content]
|
||||
if isinstance(content, dict) and isinstance(text := content.get("text"), str):
|
||||
return {**content, "text": text.replace(_EMAIL, _MASK)}
|
||||
return content
|
||||
|
||||
|
||||
def _masked_row(row: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
return {**row, "content": _masked(row["content"])} if "content" in row else dict(row)
|
||||
|
||||
|
||||
def _guardrail_peer(request: Request) -> Reply:
|
||||
verdict: Final = request.target.split("/")[1]
|
||||
rows: Final = _JSON_OBJECTS.validate_python(
|
||||
_JSON_OBJECT.validate_json(request.body).get("structured_messages") or []
|
||||
)
|
||||
match verdict:
|
||||
case "block":
|
||||
return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": _BLOCK_REASON}).encode())
|
||||
case "rewrite" if rows:
|
||||
rewritten: Final = [_masked_row(row) for row in rows]
|
||||
return Reply(body=json.dumps({"action": "GUARDRAIL_INTERVENED", "structured_messages": rewritten}).encode())
|
||||
case _:
|
||||
return Reply(body=json.dumps({"action": "NONE"}).encode())
|
||||
|
||||
|
||||
def _delete_guardrail(gateway: Gateway, identity: str) -> None:
|
||||
deleted: Final = gateway.request("DELETE", f"/guardrails/{identity}")
|
||||
assert deleted.status_code == 200, deleted.text
|
||||
|
||||
|
||||
def _guardrail(gateway: Gateway, scenario: Scenario, api_base: str, mode: str, **extra: JsonValue) -> str:
|
||||
name: Final = f"long-tool-names-{mode}-{uuid.uuid4().hex}"
|
||||
params: Final[dict[str, JsonValue]] = {
|
||||
"guardrail": "generic_guardrail_api",
|
||||
"mode": mode,
|
||||
"default_on": False,
|
||||
"api_base": api_base,
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
**extra,
|
||||
}
|
||||
created: Final = gateway.post("/guardrails", {"guardrail": {"guardrail_name": name, "litellm_params": params}})
|
||||
scenario.cleanups.callback(_delete_guardrail, gateway, string_value(created["guardrail_id"]))
|
||||
return name
|
||||
|
||||
|
||||
def _anthropic_model(scenario: Scenario, provider: Wire) -> str:
|
||||
return scenario.model(model="anthropic/claude-opus-5", api_base=provider.url, api_key="synthetic-anthropic-key")
|
||||
|
||||
|
||||
def _sse_tool_use_names(text: str) -> tuple[str, ...]:
|
||||
events: Final = tuple(
|
||||
_JSON_OBJECT.validate_json(line.removeprefix("data: "))
|
||||
for line in text.splitlines()
|
||||
if line.startswith("data: ")
|
||||
)
|
||||
starts: Final = tuple(
|
||||
object_value(event["content_block"]) for event in events if event.get("type") == "content_block_start"
|
||||
)
|
||||
return tuple(string_value(block["name"]) for block in starts if block.get("type") == "tool_use")
|
||||
|
||||
|
||||
def _answered_tool_names(response: httpx.Response, stream: bool) -> tuple[str, ...]:
|
||||
if stream:
|
||||
return _sse_tool_use_names(response.text)
|
||||
content: Final = _JSON_OBJECTS.validate_python(_JSON_OBJECT.validate_json(response.content)["content"])
|
||||
return tuple(string_value(block["name"]) for block in content if block.get("type") == "tool_use")
|
||||
|
||||
|
||||
def _vendor_tool_names(vendor_requests: tuple[Request, ...]) -> tuple[str, ...]:
|
||||
bodies: Final = tuple(_JSON_OBJECT.validate_json(request.body) for request in vendor_requests)
|
||||
tools: Final = tuple(
|
||||
object_value(tool) for body in bodies for tool in _JSON_OBJECTS.validate_python(body.get("tools") or [])
|
||||
)
|
||||
return tuple(string_value(object_value(tool["function"])["name"]) for tool in tools)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", (False, True), ids=("json", "sse"))
|
||||
@pytest.mark.parametrize("length", (64, 65, 128, 200))
|
||||
@pytest.mark.parametrize("verdict", ("echo", "rewrite"))
|
||||
@pytest.mark.parametrize("mode", ("pre_call", "during_call"))
|
||||
def test_guardrail_forwards_every_tool_name_the_caller_sent(
|
||||
gateway: Gateway, mode: str, verdict: str, length: int, stream: bool
|
||||
) -> None:
|
||||
name: Final = _tool_name(length)
|
||||
with (
|
||||
wire_server(_anthropic_peer) as provider,
|
||||
wire_server(_guardrail_peer) as vendor,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = _anthropic_model(scenario, provider)
|
||||
guardrail: Final = _guardrail(gateway, scenario, f"{vendor.url}/{verdict}", mode)
|
||||
short: Final = gateway.request("POST", "/v1/messages", _every_carrier(model, _SHORT_NAME, guardrail, stream))
|
||||
response: Final = gateway.request("POST", "/v1/messages", _every_carrier(model, name, guardrail, stream))
|
||||
sent: Final = provider.drain()
|
||||
screened: Final = vendor.drain()
|
||||
|
||||
assert (short.status_code, response.status_code) == (200, 200), (short.text, response.text)
|
||||
assert len(sent) == 2, sent
|
||||
assert len(screened) == 2, screened
|
||||
outbound: Final = _JSON_OBJECT.validate_json(sent[1].body)
|
||||
assert _carriers(outbound) == ((_SEARCH, name), name, (_SEARCH, name), (name,)), outbound
|
||||
assert outbound == _JSON_OBJECT.validate_json(sent[0].body.decode().replace(_SHORT_NAME, name)), outbound
|
||||
assert _answered_tool_names(response, stream) == (name,), response.text
|
||||
assert all(len(seen) <= _GUARDRAIL_NAME_LIMIT for seen in _vendor_tool_names(screened)), screened
|
||||
rewritten: Final = mode == "pre_call" and verdict == "rewrite"
|
||||
assert (_MASK in sent[1].body.decode(), _EMAIL in sent[1].body.decode()) == (rewritten, not rewritten), outbound
|
||||
|
||||
|
||||
@pytest.mark.parametrize("verdict", ("echo", "rewrite"))
|
||||
def test_short_tool_named_like_a_shortened_history_tool_keeps_its_name(gateway: Gateway, verdict: str) -> None:
|
||||
history_only: Final = _tool_name(90)
|
||||
alias: Final = truncate_tool_name(history_only)
|
||||
with (
|
||||
wire_server(_anthropic_peer) as provider,
|
||||
wire_server(_guardrail_peer) as vendor,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"model": _anthropic_model(scenario, provider),
|
||||
"max_tokens": 32,
|
||||
"guardrails": [_guardrail(gateway, scenario, f"{vendor.url}/{verdict}", "pre_call")],
|
||||
"tools": [_tool(alias)],
|
||||
"tool_choice": {"type": "tool", "name": alias},
|
||||
"messages": [
|
||||
{"role": "user", "content": "list my workspaces"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "tool_use", "id": "toolu_01", "name": history_only, "input": {}}],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "tool_result", "tool_use_id": "toolu_01", "content": "workspace-alpha"},
|
||||
{"type": "text", "text": f"share it with {_EMAIL}"},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
response: Final = gateway.request("POST", "/v1/messages", body)
|
||||
sent: Final = provider.drain()
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(sent) == 1, sent
|
||||
outbound: Final = _JSON_OBJECT.validate_json(sent[0].body)
|
||||
assert _carriers(outbound) == ((alias,), alias, (history_only,), ()), outbound
|
||||
assert _answered_tool_names(response, stream=False) == (alias,), response.text
|
||||
|
||||
|
||||
def test_logging_only_guardrail_forwards_the_long_tool_names_untouched(gateway: Gateway) -> None:
|
||||
name: Final = _tool_name(90)
|
||||
with (
|
||||
wire_server(_anthropic_peer) as provider,
|
||||
wire_server(_guardrail_peer) as vendor,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
body: Final = _every_carrier(
|
||||
_anthropic_model(scenario, provider),
|
||||
name,
|
||||
_guardrail(gateway, scenario, f"{vendor.url}/rewrite", "logging_only"),
|
||||
False,
|
||||
)
|
||||
response: Final = gateway.request("POST", "/v1/messages", body)
|
||||
sent: Final = provider.drain()
|
||||
eventually(vendor.received.qsize, lambda count: count >= 1)
|
||||
screened: Final = vendor.drain()
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(sent) == 1, sent
|
||||
outbound: Final = _JSON_OBJECT.validate_json(sent[0].body)
|
||||
assert _carriers(outbound) == ((_SEARCH, name), name, (_SEARCH, name), (name,)), outbound
|
||||
assert _EMAIL in sent[0].body.decode() and _MASK not in sent[0].body.decode(), outbound
|
||||
assert _answered_tool_names(response, stream=False) == (name,), response.text
|
||||
assert all(len(seen) <= _GUARDRAIL_NAME_LIMIT for seen in _vendor_tool_names(screened)), screened
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fallback", ("fail_open", "fail_closed"))
|
||||
def test_unreachable_guardrail_follows_its_fallback_with_long_tool_names(gateway: Gateway, fallback: str) -> None:
|
||||
name: Final = _tool_name(90)
|
||||
with wire_server(_anthropic_peer) as provider, gateway.scenario() as scenario:
|
||||
guardrail: Final = _guardrail(
|
||||
gateway, scenario, f"http://127.0.0.1:{allocate_port()}", "pre_call", unreachable_fallback=fallback
|
||||
)
|
||||
body: Final = _every_carrier(_anthropic_model(scenario, provider), name, guardrail, False)
|
||||
response: Final = gateway.request("POST", "/v1/messages", body)
|
||||
sent: Final = provider.drain()
|
||||
|
||||
if fallback == "fail_closed":
|
||||
assert (response.status_code, sent) == (500, ()), response.text
|
||||
return
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(sent) == 1, sent
|
||||
outbound: Final = _JSON_OBJECT.validate_json(sent[0].body)
|
||||
assert _carriers(outbound) == ((_SEARCH, name), name, (_SEARCH, name), (name,)), outbound
|
||||
|
||||
|
||||
def test_blocking_guardrail_stops_a_long_tool_name_request_before_the_provider(gateway: Gateway) -> None:
|
||||
with (
|
||||
wire_server(_anthropic_peer) as provider,
|
||||
wire_server(_guardrail_peer) as vendor,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
body: Final = _every_carrier(
|
||||
_anthropic_model(scenario, provider),
|
||||
_tool_name(90),
|
||||
_guardrail(gateway, scenario, f"{vendor.url}/block", "pre_call"),
|
||||
stream=False,
|
||||
)
|
||||
response: Final = gateway.request("POST", "/v1/messages", body)
|
||||
sent: Final = provider.drain()
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert _BLOCK_REASON in response.text, response.text
|
||||
assert sent == (), sent
|
||||
|
||||
|
||||
@pytest.mark.parametrize("length", (64, 90))
|
||||
def test_bridged_route_answers_with_the_tool_name_the_caller_sent(gateway: Gateway, length: int) -> None:
|
||||
name: Final = _tool_name(length)
|
||||
with (
|
||||
wire_server(_chat_peer) as provider,
|
||||
wire_server(_guardrail_peer) as vendor,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(
|
||||
model="mistral/mistral-small-latest", api_base=provider.url, api_key="synthetic-mistral-key"
|
||||
)
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"model": model,
|
||||
"max_tokens": 32,
|
||||
"guardrails": [_guardrail(gateway, scenario, f"{vendor.url}/echo", "pre_call")],
|
||||
"tools": [_tool(name)],
|
||||
"messages": [{"role": "user", "content": "list my workspaces"}],
|
||||
}
|
||||
response: Final = gateway.request("POST", "/v1/messages", body)
|
||||
sent: Final = provider.drain()
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(sent) == 1, sent
|
||||
assert _answered_tool_names(response, stream=False) == (name,), response.text
|
||||
|
||||
|
||||
_PARAMETERS: Final[dict[str, JsonValue]] = {"type": "object", "properties": {}}
|
||||
|
||||
|
||||
def _openai_shaped_body(endpoint: str, model: str, name: str, guardrail: str) -> dict[str, JsonValue]:
|
||||
shared: Final[dict[str, JsonValue]] = {"model": model, "guardrails": [guardrail]}
|
||||
if endpoint == "/v1/chat/completions":
|
||||
return {
|
||||
**shared,
|
||||
"messages": [{"role": "user", "content": "list my workspaces"}],
|
||||
"tools": [{"type": "function", "function": {"name": name, "parameters": _PARAMETERS}}],
|
||||
"tool_choice": {"type": "function", "function": {"name": name}},
|
||||
}
|
||||
return {
|
||||
**shared,
|
||||
"input": "list my workspaces",
|
||||
"tools": [{"type": "function", "name": name, "parameters": _PARAMETERS}],
|
||||
"tool_choice": {"type": "function", "name": name},
|
||||
}
|
||||
|
||||
|
||||
def _openai_shaped_tool_calls(endpoint: str, response: httpx.Response) -> tuple[str, ...]:
|
||||
body: Final = _JSON_OBJECT.validate_json(response.content)
|
||||
if endpoint == "/v1/chat/completions":
|
||||
message: Final = object_value(object_value(_JSON_OBJECTS.validate_python(body["choices"])[0])["message"])
|
||||
calls: Final = _JSON_OBJECTS.validate_python(message.get("tool_calls") or [])
|
||||
return tuple(string_value(object_value(call["function"])["name"]) for call in calls)
|
||||
items: Final = _JSON_OBJECTS.validate_python(body["output"])
|
||||
return tuple(string_value(item["name"]) for item in items if item.get("type") == "function_call")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", ("/v1/chat/completions", "/v1/responses"))
|
||||
def test_openai_shaped_routes_forward_the_long_tool_name_the_caller_sent(gateway: Gateway, endpoint: str) -> None:
|
||||
name: Final = _tool_name(90)
|
||||
with (
|
||||
wire_server(_anthropic_peer) as provider,
|
||||
wire_server(_guardrail_peer) as vendor,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
body: Final = _openai_shaped_body(
|
||||
endpoint,
|
||||
_anthropic_model(scenario, provider),
|
||||
name,
|
||||
_guardrail(gateway, scenario, f"{vendor.url}/echo", "pre_call"),
|
||||
)
|
||||
response: Final = gateway.request("POST", endpoint, body)
|
||||
sent: Final = provider.drain()
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(sent) == 1, sent
|
||||
assert _carriers(_JSON_OBJECT.validate_json(sent[0].body))[:2] == ((name,), name), sent
|
||||
assert _openai_shaped_tool_calls(endpoint, response) == (name,), response.text
|
||||
|
||||
|
||||
def _forced_tool_kwargs(model: str, name: str, guardrail: str) -> dict[str, object]:
|
||||
return {
|
||||
"model": model,
|
||||
"max_tokens": 32,
|
||||
"tools": [_tool(name)],
|
||||
"tool_choice": {"type": "tool", "name": name},
|
||||
"messages": [{"role": "user", "content": "list my workspaces"}],
|
||||
"extra_body": {"guardrails": [guardrail]},
|
||||
}
|
||||
|
||||
|
||||
def _origin(gateway: Gateway) -> str:
|
||||
return str(gateway.client.base_url).rstrip("/")
|
||||
|
||||
|
||||
def test_sync_anthropic_sdk_gets_the_forced_long_tool_back(gateway: Gateway) -> None:
|
||||
name: Final = _tool_name(90)
|
||||
with (
|
||||
wire_server(_anthropic_peer) as provider,
|
||||
wire_server(_guardrail_peer) as vendor,
|
||||
gateway.scenario() as scenario,
|
||||
httpx.Client(timeout=30, trust_env=False) as transport,
|
||||
anthropic.Anthropic(
|
||||
api_key=gateway.key, base_url=_origin(gateway), max_retries=0, http_client=transport
|
||||
) as client,
|
||||
):
|
||||
kwargs: Final = _forced_tool_kwargs(
|
||||
_anthropic_model(scenario, provider),
|
||||
name,
|
||||
_guardrail(gateway, scenario, f"{vendor.url}/echo", "pre_call"),
|
||||
)
|
||||
message: Final = client.messages.create(**kwargs)
|
||||
sent: Final = provider.drain()
|
||||
|
||||
assert [block.name for block in message.content if block.type == "tool_use"] == [name], message
|
||||
assert _carriers(_JSON_OBJECT.validate_json(sent[0].body))[:2] == ((name,), name), sent
|
||||
|
||||
|
||||
def test_async_anthropic_sdk_streams_the_forced_long_tool_back(gateway: Gateway) -> None:
|
||||
name: Final = _tool_name(90)
|
||||
|
||||
async def streamed(client: anthropic.AsyncAnthropic, kwargs: dict[str, object]) -> tuple[str, ...]:
|
||||
async with client.messages.stream(**kwargs) as stream:
|
||||
final: Final = await stream.get_final_message()
|
||||
return tuple(block.name for block in final.content if block.type == "tool_use")
|
||||
|
||||
async def run(kwargs: dict[str, object]) -> tuple[str, ...]:
|
||||
async with (
|
||||
httpx.AsyncClient(timeout=30, trust_env=False) as transport,
|
||||
anthropic.AsyncAnthropic(
|
||||
api_key=gateway.key, base_url=_origin(gateway), max_retries=0, http_client=transport
|
||||
) as client,
|
||||
):
|
||||
return await streamed(client, kwargs)
|
||||
|
||||
with (
|
||||
wire_server(_anthropic_peer) as provider,
|
||||
wire_server(_guardrail_peer) as vendor,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
kwargs: Final = _forced_tool_kwargs(
|
||||
_anthropic_model(scenario, provider),
|
||||
name,
|
||||
_guardrail(gateway, scenario, f"{vendor.url}/echo", "pre_call"),
|
||||
)
|
||||
answered: Final = asyncio.run(run(kwargs))
|
||||
sent: Final = provider.drain()
|
||||
|
||||
assert answered == (name,), answered
|
||||
assert _carriers(_JSON_OBJECT.validate_json(sent[0].body))[:2] == ((name,), name), sent
|
||||
|
||||
|
||||
def test_concurrent_long_tool_name_requests_each_keep_their_own_names(gateway: Gateway) -> None:
|
||||
names: Final = tuple(f"{_tool_name(85)}_{index:02d}" for index in range(24))
|
||||
with (
|
||||
wire_server(_anthropic_peer) as provider,
|
||||
wire_server(_guardrail_peer) as vendor,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = _anthropic_model(scenario, provider)
|
||||
guardrail: Final = _guardrail(gateway, scenario, f"{vendor.url}/rewrite", "pre_call")
|
||||
with ThreadPoolExecutor(max_workers=12) as pool:
|
||||
responses: Final = tuple(
|
||||
pool.map(
|
||||
lambda name: gateway.request("POST", "/v1/messages", _every_carrier(model, name, guardrail, False)),
|
||||
names,
|
||||
)
|
||||
)
|
||||
sent: Final = provider.drain()
|
||||
|
||||
assert [response.status_code for response in responses] == [200] * len(names), [r.text for r in responses]
|
||||
assert [_answered_tool_names(response, stream=False) for response in responses] == [(name,) for name in names]
|
||||
forwarded: Final = sorted(_carriers(_JSON_OBJECT.validate_json(request.body)) for request in sent)
|
||||
assert forwarded == sorted(((_SEARCH, name), name, (_SEARCH, name), (name,)) for name in names), forwarded
|
||||
|
|
@ -18,6 +18,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import Streami
|
|||
from litellm.llms.anthropic.chat.guardrail_translation.handler import (
|
||||
AnthropicMessagesHandler,
|
||||
)
|
||||
from litellm.llms.anthropic.pass_through.adapters.transformation import truncate_tool_name
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs
|
||||
|
||||
|
||||
|
|
@ -2693,6 +2694,143 @@ class TestStructuredWriteBackKeepsToolResults:
|
|||
]
|
||||
|
||||
|
||||
class NamelessComputerToolGuardrail(CustomGuardrail):
|
||||
"""Guardrail that appends an OpenAI-format computer tool without a function name."""
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
computer = {"display_width_px": 1024, "display_height_px": 768}
|
||||
inputs["tools"] = [
|
||||
*(inputs.get("tools") or []),
|
||||
{"type": "computer_20250124", "function": {"parameters": computer}},
|
||||
]
|
||||
return inputs
|
||||
|
||||
|
||||
class TestLongToolNamesSurviveTheWriteBack:
|
||||
"""Names over OpenAI's 64-char limit are shortened only for the guardrail, never in the request (LIT-9528)."""
|
||||
|
||||
_PREFIX: Final = "mcp__plugin_support-kit_remote-workspace__getAccessibleWorkspace"
|
||||
_SHORT_NAMES: Final = {"searched": "searched_tool", "called": "called_tool", "dropped": "dropped_tool"}
|
||||
_LONG_NAMES: Final = {
|
||||
"searched": _PREFIX + "Resources",
|
||||
"called": _PREFIX + "Projects",
|
||||
"dropped": _PREFIX + "Archives",
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _tool(name: str) -> dict:
|
||||
return {"name": name, "description": "Returns OK", "input_schema": {"type": "object", "properties": {}}}
|
||||
|
||||
def _request(self, searched: str, called: str, dropped: str) -> dict:
|
||||
return {
|
||||
"model": "claude-sonnet-5-5",
|
||||
"tools": [self._tool("ToolSearch"), self._tool(searched), self._tool(called)],
|
||||
"tool_choice": {"type": "tool", "name": called},
|
||||
"messages": [
|
||||
{"role": "user", "content": "find the tools for bob@example.com"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Searching."},
|
||||
{"type": "tool_use", "id": "toolu_01", "name": "ToolSearch", "input": {"query": "x"}},
|
||||
{"type": "tool_use", "id": "toolu_02", "name": called, "input": {"page": 2}},
|
||||
{"type": "tool_use", "id": "toolu_03", "name": dropped, "input": {"all": True}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_01",
|
||||
"content": [{"type": "tool_reference", "tool_name": searched}],
|
||||
},
|
||||
{"type": "tool_result", "tool_use_id": "toolu_02", "content": "OK"},
|
||||
{"type": "tool_result", "tool_use_id": "toolu_03", "content": "OK"},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "say done"},
|
||||
],
|
||||
}
|
||||
|
||||
async def _processed(self, guardrail: CustomGuardrail, names: dict) -> dict:
|
||||
data = self._request(**names)
|
||||
await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
return {key: value for key, value in data.items() if key != "metadata"}
|
||||
|
||||
def _with_long_names(self, processed: dict) -> dict:
|
||||
text = json.dumps(processed)
|
||||
for role, short in self._SHORT_NAMES.items():
|
||||
text = text.replace(f'"{short}"', f'"{self._LONG_NAMES[role]}"')
|
||||
return json.loads(text)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"guardrail",
|
||||
[
|
||||
MockPassThroughGuardrail(guardrail_name="test"),
|
||||
ToolAppendingGuardrail(guardrail_name="test"),
|
||||
MockStructuredMaskingGuardrail(),
|
||||
],
|
||||
ids=["echoed-tools", "injected-tool", "rewritten-messages"],
|
||||
)
|
||||
async def test_long_tool_names_come_back_exactly_like_short_ones(self, guardrail):
|
||||
short = await self._processed(guardrail, self._SHORT_NAMES)
|
||||
long = await self._processed(guardrail, self._LONG_NAMES)
|
||||
|
||||
assert long == self._with_long_names(short)
|
||||
assert long["tool_choice"] == {"type": "tool", "name": self._LONG_NAMES["called"]}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"guardrail",
|
||||
[MockPassThroughGuardrail(guardrail_name="test"), MockStructuredMaskingGuardrail()],
|
||||
ids=["echoed-tools", "rewritten-messages"],
|
||||
)
|
||||
async def test_short_tool_named_like_a_shortened_history_tool_keeps_its_name(self, guardrail):
|
||||
alias = truncate_tool_name(self._LONG_NAMES["dropped"])
|
||||
names = {**self._LONG_NAMES, "called": alias}
|
||||
|
||||
data = await self._processed(guardrail, names)
|
||||
|
||||
assert [tool["name"] for tool in data["tools"]] == ["ToolSearch", self._LONG_NAMES["searched"], alias]
|
||||
assert data["tool_choice"] == {"type": "tool", "name": alias}
|
||||
assert data["messages"][1]["content"][2] == {
|
||||
"type": "tool_use",
|
||||
"id": "toolu_02",
|
||||
"name": alias,
|
||||
"input": {"page": 2},
|
||||
}
|
||||
|
||||
async def test_history_tool_with_the_same_shortened_name_cannot_rename_a_declared_tool(self):
|
||||
declared, historical = (
|
||||
"mcp__plugin_support-kit_remote-workspace__getAccessibleWorkspaceResources_7aeb",
|
||||
"mcp__plugin_support-kit_remote-workspace__getAccessibleWorkspaceResources_c034",
|
||||
)
|
||||
assert truncate_tool_name(declared) == truncate_tool_name(historical)
|
||||
|
||||
data = await self._processed(
|
||||
MockPassThroughGuardrail(guardrail_name="test"),
|
||||
{**self._LONG_NAMES, "called": declared, "dropped": historical},
|
||||
)
|
||||
|
||||
assert [tool["name"] for tool in data["tools"]] == ["ToolSearch", self._LONG_NAMES["searched"], declared]
|
||||
|
||||
async def test_guardrail_tool_without_a_function_name_still_reaches_the_request(self):
|
||||
data = await self._processed(NamelessComputerToolGuardrail(guardrail_name="test"), self._LONG_NAMES)
|
||||
|
||||
assert data["tools"][-1] == {
|
||||
"type": "computer_20250124",
|
||||
"name": "computer",
|
||||
"display_width_px": 1024,
|
||||
"display_height_px": 768,
|
||||
}
|
||||
|
||||
|
||||
class TestAnthropicMessagesHandlerStreamingScanKey:
|
||||
"""get_streaming_scan_key mirrors what process_output_streaming_response would scan"""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue