diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 7e3fda06b1a..d3662fe85ba 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -53,6 +53,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, + guardrail_scans_tool_descriptions, list_field, merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, @@ -887,7 +888,11 @@ class AnthropicMessagesHandler(BaseTranslation): pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) # Step 2: Apply guardrail to all texts and tool calls in batch - if texts_to_check or tool_calls_to_check: + if ( + texts_to_check + or tool_calls_to_check + or (tools_to_check and guardrail_scans_tool_descriptions(guardrail_to_apply)) + ): inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: inputs["images"] = images_to_check diff --git a/litellm/llms/base_llm/guardrail_translation/tool_descriptions.py b/litellm/llms/base_llm/guardrail_translation/tool_descriptions.py index 30e71560cb1..2d7f20729c7 100644 --- a/litellm/llms/base_llm/guardrail_translation/tool_descriptions.py +++ b/litellm/llms/base_llm/guardrail_translation/tool_descriptions.py @@ -2,98 +2,45 @@ from __future__ import annotations import copy import functools -import itertools -from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import Final, TypeAlias, TypeVar from fastapi import HTTPException from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import TypeIs # noqa: TID251 # narrows JSON in place without copying it -DescriptionPath: TypeAlias = tuple[str | int, ...] -_Located: TypeAlias = tuple[DescriptionPath, JsonValue] +from litellm.proxy.mcp_server.utils import JSONLeafPath, json_string_leaves + +DescriptionPath: TypeAlias = JSONLeafPath _TOOLS: Final = TypeAdapter(list[JsonValue]) _Tool = TypeVar("_Tool") # rebind-ok: a TypeVar is declared by plain assignment -_TOOL_SCHEMA_KEYS: Final = ("parameters", "input_schema") -_SUBSCHEMA_KEYWORDS: Final = ( - "items", - "additionalItems", - "unevaluatedItems", - "contains", - "additionalProperties", - "unevaluatedProperties", - "propertyNames", - "not", - "if", - "then", - "else", -) -_SUBSCHEMA_LIST_KEYWORDS: Final = ("items", "prefixItems", "anyOf", "oneOf", "allOf") -_SUBSCHEMA_MAP_KEYWORDS: Final = ( - "properties", - "patternProperties", - "$defs", - "definitions", - "dependentSchemas", - "dependencies", +_SCHEMA_VALUE_KEYWORDS: Final = frozenset({"default", "examples", "enum", "const"}) +_NAMED_SUBSCHEMA_MAPS: Final = frozenset( + {"properties", "patternProperties", "$defs", "definitions", "dependentSchemas", "dependencies"} ) -def tool_description_entries(tools: Sequence[JsonValue]) -> tuple[tuple[DescriptionPath, str], ...]: - holders: Final = tuple( - itertools.chain.from_iterable(_tool_holders(tool, (index,)) for index, tool in enumerate(tools)) +def _inside_schema_value(path: DescriptionPath) -> bool: + return any( + key in _SCHEMA_VALUE_KEYWORDS and (index == 0 or path[index - 1] not in _NAMED_SUBSCHEMA_MAPS) + for index, key in enumerate(path) ) + + +def tool_description_entries(tools: Sequence[JsonValue]) -> tuple[tuple[DescriptionPath, str], ...] | None: + """Every ``description`` string in ``tools``, or None when they are nested too deeply to read.""" + leaves: Final = json_string_leaves(list(tools)) + if leaves is None: + return None return tuple( - itertools.chain( - itertools.chain.from_iterable(_own_description(holder, path) for path, holder in holders), - itertools.chain.from_iterable(_own_description(schema, path) for path, schema in _schemas(holders)), - ) + (path, text) + for path, text in leaves + if path and path[-1] == "description" and text.strip() and not _inside_schema_value(path[:-1]) ) -def _tool_holders(tool: JsonValue, path: DescriptionPath) -> Iterator[tuple[DescriptionPath, Mapping[str, JsonValue]]]: - if not isinstance(tool, dict): - return - yield path, tool - if isinstance(function := tool.get("function"), dict): - yield (*path, "function"), function - - -def _schemas( - holders: Sequence[tuple[DescriptionPath, Mapping[str, JsonValue]]], -) -> Iterator[tuple[DescriptionPath, Mapping[str, JsonValue]]]: - roots: Final = itertools.chain.from_iterable( - (((*path, key), holder.get(key)) for key in _TOOL_SCHEMA_KEYS) for path, holder in holders - ) - pending: Final[list[_Located]] = list(roots) # mutable-ok: work stack keeps the schema walk non-recursive - while pending: - path, schema = pending.pop() - if isinstance(schema, dict): - yield path, schema - pending.extend(_subschemas(schema, path)) - - -def _subschemas(schema: Mapping[str, JsonValue], path: DescriptionPath) -> Iterator[tuple[DescriptionPath, JsonValue]]: - for key in _SUBSCHEMA_KEYWORDS: - yield (*path, key), schema.get(key) - for key in _SUBSCHEMA_LIST_KEYWORDS: - if isinstance(children := schema.get(key), list): - for index, child in enumerate(children): - yield (*path, key, index), child - for key in _SUBSCHEMA_MAP_KEYWORDS: - if isinstance(named := schema.get(key), dict): - for name, child in named.items(): - yield (*path, key, name), child - - -def _own_description(node: Mapping[str, JsonValue], path: DescriptionPath) -> Iterator[tuple[DescriptionPath, str]]: - description: Final = node.get("description") - if isinstance(description, str) and description.strip(): - yield (*path, "description"), description - - def _is_object( value: object, ) -> TypeIs[dict[str, object]]: # guard-ok: JSON keys are str # mutable-ok: rewritten in place @@ -125,32 +72,38 @@ def with_replaced_descriptions( return rewritten +def _unscannable(guardrail_name: str | None) -> HTTPException: + return HTTPException( + status_code=400, + detail={ + "error": "tool_descriptions_unscannable", + "message": ( + f"Guardrail '{guardrail_name}' has scan_tool_descriptions_as_text on, but tools[] is not " + "plain JSON or is nested too deeply to read, so the request was rejected rather than sent " + "with unscanned tool descriptions" + ), + }, + ) + + async def scanned_tool_descriptions( tools: Sequence[_Tool], scan_text: Callable[[str], Awaitable[str]], guardrail_name: str | None, ) -> list[_Tool] | None: # mutable-ok: GenericGuardrailAPIInputs.tools is declared as a list - """Run ``scan_text`` over every tool and schema ``description`` in OpenAI-format ``tools``. + """Run ``scan_text`` over every ``description`` string in ``tools``, skipping schema values like ``default``. Returns the tools with changed descriptions written back at their original paths, or None when nothing changed. ``scan_text`` blocks by raising. """ try: - validated: Final = _TOOLS.validate_python(tools) + entries: Final = tool_description_entries(_TOOLS.validate_python(tools)) except ValidationError as error: - raise HTTPException( - status_code=400, - detail={ - "error": "tool_descriptions_unscannable", - "message": ( - f"Guardrail '{guardrail_name}' has scan_tool_descriptions_as_text on, but tools[] is not " - "plain JSON or is nested too deeply to read, so the request was rejected rather than sent " - "with unscanned tool descriptions" - ), - }, - ) from error + raise _unscannable(guardrail_name) from error + if entries is None: + raise _unscannable(guardrail_name) replacements: Final[dict[DescriptionPath, str]] = {} # mutable-ok: filled one scan at a time - for path, text in tool_description_entries(validated): + for path, text in entries: scanned = await scan_text(text) if scanned != text: replacements[path] = scanned diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index fb886e42780..9eb7d43e362 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -248,6 +248,10 @@ def effective_scan_only_tool_results_for_guardrail(guardrail_to_apply: object) - return getattr(guardrail_to_apply, "scan_only_tool_results", None) is True +def guardrail_scans_tool_descriptions(guardrail_to_apply: object) -> bool: + return getattr(guardrail_to_apply, "scan_tool_descriptions_as_text", None) is True + + def role_out_of_guardrail_scope( role: str, *, diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py index d49d5cc4c85..4cf57dbfb49 100644 --- a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py +++ b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py @@ -28,6 +28,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( ) from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, + guardrail_scans_tool_descriptions, unappliable_request_rewrite, ) from litellm.llms.custom_httpx.sse_lines import split_sse_lines @@ -245,7 +246,10 @@ class GoogleGenAIGenerateContentHandler(BaseTranslation): litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> object: text_parts: Final = _request_text_parts(data) - if not text_parts: + declarations: Final = ( + () if effective_scan_only_tool_results_for_guardrail(guardrail_to_apply) else _request_declarations(data) # pyright: ignore[reportUnknownArgumentType] # data is typed as a bare dict by BaseTranslation + ) + if not text_parts and not (declarations and guardrail_scans_tool_descriptions(guardrail_to_apply)): verbose_proxy_logger.debug("Google GenAI guardrail: no request text found, skipping") return data model: Final = data.get("model") @@ -254,9 +258,6 @@ class GoogleGenAIGenerateContentHandler(BaseTranslation): if isinstance(model, str) else GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts))) ) - declarations: Final = ( - () if effective_scan_only_tool_results_for_guardrail(guardrail_to_apply) else _request_declarations(data) # pyright: ignore[reportUnknownArgumentType] # data is typed as a bare dict by BaseTranslation - ) if declarations: inputs["tools"] = [declaration.tool for declaration in declarations] guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 5c5204ef2f4..ad93f1cdb24 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -39,6 +39,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, + guardrail_scans_tool_descriptions, list_field, merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, @@ -139,8 +140,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation): scan_only_tool_results=scan_only_tool_results, ) + tools: Final = data.get("tools") # Step 2: Apply guardrail to all texts and tool calls in batch - if texts_to_check or tool_calls_to_check: + if ( + texts_to_check + or tool_calls_to_check + or ( + isinstance(tools, list) + and tools + and not scan_only_tool_results + and guardrail_scans_tool_descriptions(guardrail_to_apply) + ) + ): inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: inputs["images"] = images_to_check @@ -156,7 +167,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if structured_messages: inputs["structured_messages"] = [structured_messages[index] for index in scoped_message_indices] # Pass tools (function definitions) to the guardrail - tools: Final = data.get("tools") if isinstance(tools, list) and tools and not scan_only_tool_results: inputs["tools"] = tools # Include model information if available diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index a18cbe13925..e9077041c38 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -63,6 +63,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, + guardrail_scans_tool_descriptions, list_field, merge_guardrailed_scoped_messages, role_out_of_guardrail_scope, @@ -810,7 +811,9 @@ class OpenAIResponsesHandler(BaseTranslation): skip_tool=skip_tool, scan_only=scan_only, ) - if not extracted.inputs.get("texts"): + if not extracted.inputs.get("texts") and not ( + extracted.inputs.get("tools") and guardrail_scans_tool_descriptions(guardrail_to_apply) + ): return data if scoped_structured_messages: extracted.inputs["structured_messages"] = scoped_structured_messages diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index e529cba4edf..fbf5dd1421c 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -70,19 +70,41 @@ def _function_fields(tool: Tool) -> Tool: return function if function is not None else MappingProxyType({}) -def _member_description(key: str, value: object, prefix: str, suffix: str) -> object: +def _member_description(key: str, value: object, prefixes: Sequence[str], suffix: str) -> object: if key != "description" or not isinstance(value, str): return value - return value.replace(prefix, "", 1).replace(suffix, "", 1) + prefix: Final = next((prefix for prefix in prefixes if value.startswith(prefix)), "") + return value.removeprefix(prefix).replace(suffix, "", 1) -def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_description: str) -> Tool: +def _namespace_prefix(namespace_description: str) -> str: + return f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}" if namespace_description else "" + + +def _guardrailed_namespace_description(namespace_description: str, guardrailed_members: Iterable[Tool]) -> str: + """The namespace description as the guardrail returned it inside its flattened members' descriptions.""" + if not namespace_description: + return namespace_description + parts: Final = namespace_description.count(NAMESPACE_DESCRIPTION_SEPARATOR) + 1 + descriptions: Final = (_function_fields(member).get("description") for member in guardrailed_members) + return next( + ( + NAMESPACE_DESCRIPTION_SEPARATOR.join(description.split(NAMESPACE_DESCRIPTION_SEPARATOR)[:parts]) + for description in descriptions + if isinstance(description, str) + and not description.startswith(_namespace_prefix(namespace_description)) + and description.count(NAMESPACE_DESCRIPTION_SEPARATOR) >= parts + ), + namespace_description, + ) + + +def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, prefixes: Sequence[str]) -> Tool: flattened_function: Final = _function_fields(flattened) - prefix: Final = f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}" if namespace_description else "" suffix: Final = custom_tool_grammar_suffix(member.get("format")) if member.get("type") == "custom" else "" changed_function: Final = MappingProxyType( { - key: _member_description(key, value, prefix, suffix) + key: _member_description(key, value, prefixes, suffix) for key, value in _function_fields(guardrailed).items() if flattened_function.get(key) != value } @@ -102,14 +124,14 @@ def _rebuilt_flattened_members( flattened_group: Sequence[Tool], group_keys: Sequence[IndexedKey], guardrailed_by_key: Mapping[IndexedKey, Tool], - namespace_description: str, + prefixes: Sequence[str], ) -> tuple[Tool | None, ...]: return tuple( None if key not in guardrailed_by_key else member if guardrailed_by_key[key] == flattened - else _rebuilt_member(member, flattened, guardrailed_by_key[key], namespace_description) + else _rebuilt_member(member, flattened, guardrailed_by_key[key], prefixes) for member, flattened, key in zip(flattened_members, flattened_group, group_keys) ) @@ -122,13 +144,17 @@ def _rebuilt_namespace( guardrailed_by_key: Mapping[IndexedKey, Tool], ) -> tuple[Tool, ...]: namespace_description: Final = str(original.get("description") or "") + guardrailed_description: Final = _guardrailed_namespace_description( + namespace_description, (guardrailed_by_key[key] for key in group_keys if key in guardrailed_by_key) + ) + prefixes: Final = (_namespace_prefix(namespace_description), _namespace_prefix(guardrailed_description)) rebuilt_flattened: Final = iter( _rebuilt_flattened_members( tuple(member for member in members if _has_chat_tool(member)), flattened_group, group_keys, guardrailed_by_key, - namespace_description, + prefixes, ) ) rebuilt_members: Final = tuple( @@ -138,7 +164,10 @@ def _rebuilt_namespace( ) if not rebuilt_members: return () - return ({**original, "tools": list(rebuilt_members)},) + description_update: Final = ( + {"description": guardrailed_description} if guardrailed_description != namespace_description else {} + ) + return ({**original, **description_update, "tools": list(rebuilt_members)},) def _merged_original( diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 3c00b282181..e5d71794f8d 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -3363,6 +3363,28 @@ _TOOL_DEFINITION_REQUESTS: Final = ( }, id="responses_namespace", ), + pytest.param( + CallTypes.aresponses.value, + { + "input": "weather in Paris?", + "tools": [ + { + "type": "namespace", + "name": "weather", + "description": _POISONED_DESCRIPTION, + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": "Gets the weather.", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ], + } + ], + }, + id="responses_namespace_description", + ), pytest.param( CallTypes.anthropic_messages.value, { @@ -3415,6 +3437,66 @@ _GEMINI_TOOL_DEFINITION_REQUESTS: Final = ( _TOOL_DESCRIPTION_SCAN_REQUESTS: Final = (*_TOOL_DEFINITION_REQUESTS, *_GEMINI_TOOL_DEFINITION_REQUESTS) +_IMAGE_URL: Final = "https://example.com/paris.png" +_POISONED_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "description": _POISONED_DESCRIPTION, "parameters": _WEATHER_SCHEMA}, +} +_REQUESTS_WITHOUT_MESSAGE_TEXT: Final = ( + pytest.param( + CallTypes.acompletion.value, + { + "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": _IMAGE_URL}}]}], + "tools": [_POISONED_FUNCTION_TOOL], + }, + id="chat_completions_image_only", + ), + pytest.param( + CallTypes.aresponses.value, + { + "input": [{"role": "user", "content": [{"type": "input_image", "image_url": _IMAGE_URL}]}], + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": _POISONED_DESCRIPTION, + "parameters": _WEATHER_SCHEMA, + } + ], + }, + id="responses_image_only", + ), + pytest.param( + CallTypes.anthropic_messages.value, + { + "max_tokens": 64, + "messages": [ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": _IMAGE_URL}}], + } + ], + "tools": [{"name": "get_weather", "description": _POISONED_DESCRIPTION, "input_schema": _WEATHER_SCHEMA}], + }, + id="anthropic_messages_image_only", + ), + pytest.param( + CallTypes.agenerate_content.value, + { + "contents": [{"role": "user", "parts": [{"inlineData": {"mimeType": "image/png", "data": "iVBORw0K"}}]}], + "tools": [ + { + "functionDeclarations": [ + {"name": "get_weather", "description": _POISONED_DESCRIPTION, "parameters": _WEATHER_SCHEMA} + ] + } + ], + }, + id="gemini_generate_content_image_only", + ), +) + + _MakeGuardrail: TypeAlias = Callable[..., CustomGuardrail] @@ -3501,6 +3583,23 @@ async def test_pre_call_blocks_request_whose_only_offending_text_is_a_tool_descr assert exc_info.value.status_code == 400 +@pytest.mark.usefixtures("_real_translation_mappings") +@pytest.mark.parametrize("make_guardrail", _TEXT_GUARDRAILS) +@pytest.mark.parametrize(("call_type", "body"), _REQUESTS_WITHOUT_MESSAGE_TEXT) +@pytest.mark.asyncio +async def test_pre_call_scans_tool_descriptions_on_requests_without_message_text( + monkeypatch: pytest.MonkeyPatch, make_guardrail: _MakeGuardrail, call_type: str, body: Mapping[str, object] +) -> None: + async def blocked(scan: object) -> bool: + try: + await _pre_call(make_guardrail(monkeypatch, mask=False, scan=scan), call_type, body) + except HTTPException: + return True + return False + + assert (await blocked(False), await blocked(True)) == (False, True) + + def _wrapped(inner: list[object], _: int) -> list[object]: return [inner] @@ -3566,8 +3665,8 @@ async def test_pre_call_masks_tool_descriptions_without_touching_the_schema( # Compare description paths and text: the Anthropic and Responses tool round-trips (unchanged from main) # add defaults such as type="custom" and strict=False to every tool a guardrail hands back. - assert tool_description_entries(_JSON_TOOLS.validate_python(result["tools"])) == tool_description_entries( - _JSON_TOOLS.validate_python(_masked(body["tools"])) + assert dict(tool_description_entries(_JSON_TOOLS.validate_python(result["tools"])) or ()) == dict( + tool_description_entries(_JSON_TOOLS.validate_python(_masked(body["tools"]))) or () ) assert "IGNORE PRIOR INSTRUCTIONS" not in json.dumps(result["tools"])