mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
fix(guardrails)!: check Anthropic and Responses request tools in tool_permission pre_call (#41011)
* fix(guardrails): check Anthropic-format request tools in tool_permission pre_call
The pre_call hook only collected tools shaped as {"type": "function", "function": {...}},
so flat tools from /v1/messages ({"name", "input_schema"}) and /v1/responses
({"type": "function", "name"}) were never checked and deny rules silently passed.
Reuse the shared anthropic_tool_names helper so every request tool name is checked,
and handle Anthropic tool_choice {"type": "tool"} in both block and rewrite modes.
Fixes #40583
* fix(guardrails): detect Responses API named tool_choice in tool_permission
Responses API named choices put the name at the top level ({"type": "function", "name": ...}),
so rewrite mode removed the denied tool but left a tool_choice pointing at it.
* test(guardrails): run tool_permission pre_call tests through real default_on gating
Build the guardrail as an always-on pre_call guardrail instead of patching should_run_guardrail.
* test(guardrails): type the new tool_permission pre_call test parameters
* fix(guardrails): keep the Responses API custom tool type distinct
Anthropic's user-defined tools carry type "custom", so pre_call mapped
that to "function" for tool_type rules. The Responses API has its own
custom tool type, and mapping it let a custom tool slip past a
tool_type deny rule when the default action is allow.
Normalize only for the formats where custom means a plain function, and
read tool names directly in the rewrite path, which never needed a type.
* test(guardrails): cover flat tools through Messages and Responses
Add blocked and permitted flat-tool cases through the running proxy, with
endpoint coverage and metadata for the current e2e harness
Require a guardrail policy verdict when checking denied tools so an
unrelated provider error cannot satisfy the regression test
Co-authored-by: Christopher Krah <christopher.krah@plexito.de>
* fix(guardrails): check namespace and nameless tools and relax tool_choice in tool permission rewrite
* fix(guardrails): check forced, grouped, loaded, and hosted tools by their real type
* fix(guardrails): drop references to removed tools and count loaded tools in tool permission rewrite
* fix(guardrails): check tools sent outside top-level tools in tool permission
* fix(guardrails): judge namespace members by their bare name as a fallback and keep rewrites valid for Gemini, Bedrock and hosted tools
A namespace member is checked by namespace__name first, and by its bare name when no rule names the joined one. Rewrite mode now prunes Gemini allowedFunctionNames and turns off forced calling when no function is left, drops a removed hosted tool from allowed_tools, reads Bedrock toolSpec and systemTool blocks by name, skips cachePoint markers, and leaves a tools field that is not a list for the provider to reject instead of failing with a 500
* fix(guardrails): let only name rules decide a namespace member's bare-name fallback
* fix(guardrails): judge malformed tool types, empty Gemini name lists and pass-through Anthropic tools like the endpoints do
A tool or tool_choice whose type is not a string now reads as untyped instead of crashing the guardrail with a 500, and a null type key is never read as a Gemini tool part. An empty allowedFunctionNames list no longer turns off the functions a rewrite kept. An Anthropic custom tool reads as a function by its input_schema on every route, so a pass-through route judges it the same as /v1/messages. A Responses WebSocket frame now forwards the extra_body tools a rewrite stripped
* fix(guardrails): judge a named hosted tool by its type and keep a nested chat custom tool custom
* fix(guardrails): keep judging Anthropic client tools and type-named server tools by their name
* fix(guardrails): judge a named tool by its name alone only when the name is its type's own
* fix(guardrails): judge the function a tool names beside a declaration list or a top-level name
* fix(guardrails): judge Gemini hosted keys whatever the tool type and a namespace's function
* fix(guardrails): judge legacy functions entries like tools and send Gemini a namespace's members
* fix(guardrails): judge a namespace member by every name a provider reads and send Gemini its schema converted
* fix(guardrails): judge a legacy functions entry as a function and relax a tool_choice in extra_body
* fix(gemini): keep a namespace whole when inlining its members would repeat a function name
* fix(guardrails): judge a namespace member by every bare name that shares its qualified name
* fix(vertex_ai): declare a repeated namespace member name once instead of sending namespace names to Gemini
* fix(guardrails): judge a tool_choice sent in extra_body like a top-level one
A forced name only extra_body.tool_choice carried was never checked, so block mode let a denied name through to the provider
* fix(guardrails): judge a Gemini namespace member returned by its bare name and force a named namespace's members
* test: keep formatter churn out of the tool_permission and Gemini test files
* fix(guardrails): judge extra_body.input tools, argument-aware alias rules, and exact forced Gemini functions
Tools loaded in extra_body.input are now checked and stripped like top-level input. A returned
namespace member is judged by the alias whose first name rule its arguments satisfy. A forced
Gemini name that matches a top-level function keeps that function. Alias grouping and Gemini
declaration dedupe now run in one pass instead of quadratic scans.
* fix(guardrails): judge and rewrite every tool field sent inside extra_body
extra_body is merged into the upstream body, so functions, function_call,
mcp_servers and toolConfig nested there reached the model unjudged. Treat
extra_body as a second request layer that is collected and rewritten the
same way as the top-level request.
* fix(guardrails): judge Gemini hosted keys by key and count only the tools a provider receives after extra_body
* fix(guardrails): relax a tool_choice that names a dropped namespace member by its bare name
Gemini and Anthropic chat send namespace members under their bare names, so rewrite mode now treats a dropped member's bare name, and an emptied namespace's own name, as removed when no kept tool still carries it
* fix(gemini,guardrails): keep a top-level function over a same-named namespace member and leave malformed tool values to the provider
Gemini and Vertex chat now declare a top-level function over a namespace member that shares its name, so a forced call gets the top-level schema. post_call judges a returned name a top-level tool carries as that tool, and pre_call no longer raises on an unhashable or non-string tool name, input item type, or function_call name.
* fix(guardrails): judge a returned namespace member named like a hosted MCP server as that member
* fix(gemini,guardrails): send a namespace as main did when another tool has a member's name
Gemini chat inlines a namespace's members only when no other tool, namespace
member, or legacy functions entry declares the same name, and otherwise sends the
namespace by its own name as main did. A tool_choice naming such a namespace
forces the namespace itself. The tool permission post_call reads top-level names
the way Gemini declares them and judges a call the Responses bridge tags with a
namespace as that member.
* fix(gemini): send colliding namespace members as namespace__name functions and judge them in post_call
When a namespace member shares a name with another tool, Gemini chat now gets that namespace's members as namespace__name functions instead of one function named after the namespace, so every function it can call is a name the rules judged. Declared names are counted once, and a forced member maps to each holder's qualified name.
post_call judges a returned namespace__name as that member, by its qualified or bare name the way pre_call does, so a bare-name allowlist or denylist applies to it.
* fix(gemini): keep generated member names unique and let a top-level function own its exact name
A namespace that listed a member twice, or a generated namespace__name that
repeated a top-level function, sent Gemini two declarations with one name, and
Gemini answered 400 where main answered 200. A name a namespace lists twice is
now one member, and a top-level function keeps its exact name.
post_call now judges a returned name a top-level function declares as that
function, never as the namespace member the name could also mean.
* fix(gemini): qualify namespace members by the namespace's top-level name
The guardrail judges a member as `name__member` from the namespace's top-level
name. Gemini built the name from `function.name` when the tool also carried one,
so a request could pass pre-call and reach Gemini under a name no rule judged.
* fix(vertex_ai): force only the namespace members Gemini is sent
A forced namespace or member on Gemini chat could still allow a generated name that a top-level function or an earlier namespace owns, so tool_choice forced an unrelated function. The forced names now come from the same collision checks the declarations use
* fix(vertex_ai): reject a forced namespace Gemini gets no member of, and resolve forced names in one pass
* fix(guardrails): judge a namespace member by a rule on its namespace's name
---------
Co-authored-by: Christopher Krah <christopher.krah@plexito.de>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Meryem Sakin
GitHub
parent
574f06594b
commit
4a7fb30b0a
13 changed files with 3627 additions and 135 deletions
|
|
@ -10,7 +10,7 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params, normalize_drop_params
|
||||
|
||||
ROUTES_WITH_TOOLS: Final = frozenset(
|
||||
{
|
||||
|
|
@ -58,6 +58,19 @@ _TOOL_ENTRY: Final = TypeAdapter(_ToolEntry)
|
|||
_WEB_SEARCH_OPTIONS: Final = TypeAdapter(dict[str, object] | None, config=ConfigDict(strict=True))
|
||||
|
||||
|
||||
def declared_function_name(tool: object) -> str | None:
|
||||
fields: Final = as_str_mapping(tool)
|
||||
if fields is None:
|
||||
return None
|
||||
function: Final = as_str_mapping(fields.get("function"))
|
||||
name: Final = (fields if function is None else function).get("name")
|
||||
return name if isinstance(name, str) and name else None
|
||||
|
||||
|
||||
def qualified_member_name(namespace: str, name: str) -> str:
|
||||
return f"{namespace}__{name}"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DropPolicy:
|
||||
drop_params: bool
|
||||
|
|
|
|||
|
|
@ -5,9 +5,11 @@ import json
|
|||
import re
|
||||
import threading
|
||||
import time
|
||||
from collections import Counter
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from copy import deepcopy
|
||||
from functools import lru_cache, partial
|
||||
from itertools import chain, repeat
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, Union, cast, get_args
|
||||
|
||||
|
|
@ -34,7 +36,11 @@ from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAcc
|
|||
from litellm.litellm_core_utils.model_param_helper import embeds_thought_signature_in_tool_call_id
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import encode_tool_call_id_with_signature
|
||||
from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning
|
||||
from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp
|
||||
from litellm.litellm_core_utils.tool_fields import (
|
||||
declared_function_name,
|
||||
qualified_member_name,
|
||||
tool_params_without_hosted_mcp,
|
||||
)
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
|
|
@ -544,6 +550,138 @@ def mapped_gemini_speech_config(
|
|||
}
|
||||
|
||||
|
||||
def _as_chat_tool(member: JsonValue) -> JsonValue:
|
||||
if not isinstance(member, dict) or "function" in member or "name" not in member:
|
||||
return member
|
||||
return {
|
||||
"type": member.get("type", "function"),
|
||||
"function": {key: value for key, value in member.items() if key != "type"},
|
||||
}
|
||||
|
||||
|
||||
def _namespace_members_or_tool(tool: JsonValue) -> tuple[JsonValue, ...]:
|
||||
if not isinstance(tool, dict) or tool.get("type") != "namespace":
|
||||
return (tool,)
|
||||
members: Final = tool.get("tools")
|
||||
return _members_named_once(tuple(map(_as_chat_tool, members))) if isinstance(members, list) else ()
|
||||
|
||||
|
||||
def _members_named_once(
|
||||
members: tuple[JsonValue, ...], repeatable: frozenset[str] = frozenset()
|
||||
) -> tuple[JsonValue, ...]:
|
||||
names: Final = tuple(map(declared_function_name, members))
|
||||
first_index: Final = {name: index for index, name in reversed(tuple(enumerate(names)))}
|
||||
return tuple(
|
||||
member
|
||||
for index, (member, name) in enumerate(zip(members, names, strict=True))
|
||||
if name is None or name in repeatable or first_index[name] == index
|
||||
)
|
||||
|
||||
|
||||
def _is_namespace(tool: JsonValue) -> bool:
|
||||
return isinstance(tool, dict) and tool.get("type") == "namespace"
|
||||
|
||||
|
||||
def _top_level_name(tool: JsonValue) -> str | None:
|
||||
name: Final = tool.get("name") if isinstance(tool, dict) else None
|
||||
return name if isinstance(name, str) and name else None
|
||||
|
||||
|
||||
def _listed(tools: JsonValue) -> tuple[JsonValue, ...]:
|
||||
return tuple(tools) if isinstance(tools, list) else ()
|
||||
|
||||
|
||||
def _declarable_names(tool: JsonValue) -> tuple[str | None, ...]:
|
||||
members: Final = _namespace_members_or_tool(tool) if _is_namespace(tool) else ()
|
||||
return (declared_function_name(tool), *map(declared_function_name, members))
|
||||
|
||||
|
||||
def _declarable_names_in(tools: tuple[JsonValue, ...]) -> tuple[str | None, ...]:
|
||||
return tuple(chain.from_iterable(map(_declarable_names, tools)))
|
||||
|
||||
|
||||
def _inlines_members(tool: JsonValue, declarable: Counter[str | None]) -> bool:
|
||||
members: Final = _namespace_members_or_tool(tool)
|
||||
return _is_namespace(tool) and all(
|
||||
name is None or declarable[name] == 1 for name in map(declared_function_name, members)
|
||||
)
|
||||
|
||||
|
||||
def _qualified_member(namespace: str, member: JsonValue) -> JsonValue:
|
||||
function: Final = member.get("function") if isinstance(member, dict) else None
|
||||
name: Final = declared_function_name(member)
|
||||
if not isinstance(member, dict) or not isinstance(function, dict) or name is None:
|
||||
return member
|
||||
return {**member, "function": {**function, "name": qualified_member_name(namespace, name)}}
|
||||
|
||||
|
||||
def _as_declared(
|
||||
tool: JsonValue, declarable: Counter[str | None], top_level_names: frozenset[str]
|
||||
) -> tuple[JsonValue, ...]:
|
||||
if _inlines_members(tool, declarable):
|
||||
return _namespace_members_or_tool(tool)
|
||||
namespace: Final = _top_level_name(tool)
|
||||
if not _is_namespace(tool) or namespace is None:
|
||||
return (tool,)
|
||||
qualified: Final = (_qualified_member(namespace, member) for member in _namespace_members_or_tool(tool))
|
||||
return tuple(member for member in qualified if declared_function_name(member) not in top_level_names)
|
||||
|
||||
|
||||
def _top_level_names_in(tools: tuple[JsonValue, ...]) -> frozenset[str]:
|
||||
top_level_tools: Final = (tool for tool in tools if not _is_namespace(tool))
|
||||
return frozenset(name for name in map(declared_function_name, top_level_tools) if name is not None)
|
||||
|
||||
|
||||
def _with_namespace_members_inlined(tools: JsonValue, outside_declarations: JsonValue = None) -> JsonValue:
|
||||
if not isinstance(tools, list) or not any(map(_is_namespace, tools)):
|
||||
return tools
|
||||
candidates: Final = (*tools, *_listed(outside_declarations))
|
||||
declarable: Final = Counter(_declarable_names_in(candidates))
|
||||
top_level_names: Final = _top_level_names_in(candidates)
|
||||
declared: Final = chain.from_iterable(_as_declared(tool, declarable, top_level_names) for tool in tools)
|
||||
return list(_members_named_once(tuple(declared), top_level_names))
|
||||
|
||||
|
||||
def _names_each_tool_declares(candidates: tuple[JsonValue, ...]) -> tuple[tuple[str, ...], ...]:
|
||||
declarable: Final = Counter(_declarable_names_in(candidates))
|
||||
top_level_names: Final = _top_level_names_in(candidates)
|
||||
declared: Final = tuple(
|
||||
tuple(filter(None, map(declared_function_name, _as_declared(tool, declarable, top_level_names))))
|
||||
for tool in candidates
|
||||
)
|
||||
owned_names: Final = chain.from_iterable(zip(repeat(index), names) for index, names in enumerate(declared))
|
||||
first_owner: Final = {name: index for index, name in reversed(tuple(owned_names))}
|
||||
return tuple(tuple(name for name in names if first_owner[name] == index) for index, names in enumerate(declared))
|
||||
|
||||
|
||||
def _qualified_names_of(
|
||||
member_name: str, tool: JsonValue, declarable: Counter[str | None], declared: tuple[str, ...]
|
||||
) -> tuple[str, ...]:
|
||||
namespace: Final = _top_level_name(tool)
|
||||
if not _is_namespace(tool) or namespace is None or _inlines_members(tool, declarable):
|
||||
return ()
|
||||
qualified: Final = qualified_member_name(namespace, member_name)
|
||||
return (qualified,) if qualified in declared else ()
|
||||
|
||||
|
||||
def _forced_function_names(name: str, tools: JsonValue) -> tuple[str, ...]:
|
||||
candidates: Final = _listed(tools)
|
||||
if any(not _is_namespace(tool) and declared_function_name(tool) == name for tool in candidates):
|
||||
return (name,)
|
||||
declarable: Final = Counter(_declarable_names_in(candidates))
|
||||
declared_by_tool: Final = tuple(zip(candidates, _names_each_tool_declares(candidates), strict=True))
|
||||
named: Final = next(
|
||||
(declared for tool, declared in declared_by_tool if _is_namespace(tool) and _top_level_name(tool) == name),
|
||||
None,
|
||||
)
|
||||
if named is not None:
|
||||
return named
|
||||
qualified: Final = chain.from_iterable(
|
||||
_qualified_names_of(name, tool, declarable, declared) for tool, declared in declared_by_tool
|
||||
)
|
||||
return tuple(qualified) or (name,)
|
||||
|
||||
|
||||
class VertexAIBaseConfig:
|
||||
def get_mapped_special_auth_params(self) -> dict:
|
||||
"""
|
||||
|
|
@ -815,7 +953,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
|
||||
return supported_params
|
||||
|
||||
def map_tool_choice_values(self, model: str, tool_choice: str | dict) -> ToolConfig | None:
|
||||
def map_tool_choice_values(
|
||||
self, model: str, tool_choice: str | dict, tools: JsonValue = None, drop_params: bool = False
|
||||
) -> ToolConfig | None:
|
||||
if tool_choice == "none":
|
||||
return ToolConfig(functionCallingConfig=FunctionCallingConfig(mode="NONE"))
|
||||
elif tool_choice == "required":
|
||||
|
|
@ -825,7 +965,17 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
elif isinstance(tool_choice, dict):
|
||||
# only supported for anthropic + mistral models - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
|
||||
name: Final = tool_choice.get("function", {}).get("name", "")
|
||||
return ToolConfig(functionCallingConfig=FunctionCallingConfig(mode="ANY", allowed_function_names=[name]))
|
||||
forced: Final = _forced_function_names(name, tools)
|
||||
if forced:
|
||||
return ToolConfig(
|
||||
functionCallingConfig=FunctionCallingConfig(mode="ANY", allowed_function_names=list(forced))
|
||||
)
|
||||
if drops_unsupported_params(drop_params):
|
||||
return None
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message=f"VertexAI can't force tool_choice '{name}', since none of that namespace's members can be sent to Gemini as a function. Name one of its functions instead, or set `litellm.drop_params = True` to drop tool_choice.",
|
||||
status_code=400,
|
||||
)
|
||||
else:
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message=f"VertexAI doesn't support tool_choice={tool_choice}. Supported tool_choice values=['auto', 'required', json object]. To drop it from the call, set `litellm.drop_params = True.",
|
||||
|
|
@ -1068,7 +1218,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
|
||||
return googleSearch, googleSearchRetrieval, enterpriseWebSearch, urlContext
|
||||
|
||||
def _map_function(self, value: list[dict], optional_params: dict, send_integer_enums: bool = False) -> list[Tools]:
|
||||
def _map_function(
|
||||
self,
|
||||
value: list[dict],
|
||||
optional_params: dict,
|
||||
send_integer_enums: bool = False,
|
||||
outside_declarations: JsonValue = None,
|
||||
) -> list[Tools]:
|
||||
"""
|
||||
Map OpenAI-style tools/functions to Vertex AI format.
|
||||
|
||||
|
|
@ -1099,7 +1255,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
# remove 'strict' from tools
|
||||
remove_strict_from_schema(schema_value)
|
||||
|
||||
for tool in cast(list[dict], schema_value): # cast-ok: a deep copy keeps the caller's tool shape
|
||||
inlined_tools: Final = _with_namespace_members_inlined(schema_value, outside_declarations)
|
||||
for tool in cast(list[dict], inlined_tools): # cast-ok: a deep copy keeps the caller's tool shape
|
||||
openai_function_object: ChatCompletionToolParamFunctionChunk | None = None
|
||||
if "function" in tool: # tools list
|
||||
_openai_function_object = ChatCompletionToolParamFunctionChunk(**tool["function"])
|
||||
|
|
@ -1734,7 +1891,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
elif (param == "tools" or param == "functions") and isinstance(value, list) and value:
|
||||
# Pass optional_params so _map_function can add toolConfig if needed
|
||||
mapped_tools = with_integer_enums_that_fit(
|
||||
partial(self._map_function, value, optional_params),
|
||||
partial(
|
||||
self._map_function,
|
||||
value,
|
||||
optional_params,
|
||||
outside_declarations=request_params.get("functions" if param == "tools" else "tools"),
|
||||
),
|
||||
allow_integer_enums and self._is_gemini_3_or_newer(model),
|
||||
)
|
||||
optional_params = self._add_tools_to_optional_params(optional_params, mapped_tools)
|
||||
|
|
@ -1742,6 +1904,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
_tool_choice_value = self.map_tool_choice_values(
|
||||
model=model,
|
||||
tool_choice=value,
|
||||
tools=[*_listed(non_default_params.get("tools")), *_listed(non_default_params.get("functions"))],
|
||||
drop_params=drop_params,
|
||||
)
|
||||
if _tool_choice_value is not None:
|
||||
optional_params["tool_choice"] = _tool_choice_value
|
||||
|
|
|
|||
|
|
@ -20233,7 +20233,7 @@
|
|||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Regex pattern applied to the tool's function name",
|
||||
"description": "Regex pattern applied to the tool's name, or to its type when the tool has no name",
|
||||
"title": "Tool Name"
|
||||
},
|
||||
"tool_type": {
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -2162,7 +2162,7 @@ _WS_FRAME: Final = TypeAdapter(dict[str, object])
|
|||
_GUARDRAIL_INFORMATION: Final = TypeAdapter(list[dict[str, object]])
|
||||
_PII_TOKENS: Final = TypeAdapter(dict[str, str])
|
||||
_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information"
|
||||
_GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools")
|
||||
_GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools", "tool_choice", "extra_body")
|
||||
_CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata")
|
||||
_CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies")
|
||||
_METADATA_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config")
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ class ToolPermissionRule(LiteLLMBaseModel):
|
|||
id: str = Field(description="Unique identifier for the rule")
|
||||
tool_name: str | None = Field(
|
||||
default=None,
|
||||
description="Regex pattern applied to the tool's function name",
|
||||
description="Regex pattern applied to the tool's name, or to its type when the tool has no name",
|
||||
)
|
||||
tool_type: str | None = Field(
|
||||
default=None,
|
||||
|
|
|
|||
|
|
@ -20,8 +20,8 @@
|
|||
- {id: guardrail.ibm_guardrails.post_call.blocks, module: guardrail, tier: P1, hook_point: post_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/ibm_guardrails/ibm_detector.py", rationale: "Output policy validation"}
|
||||
- {id: guardrail.semantic_guard.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/semantic_guard", rationale: "Semantic policy compliance"}
|
||||
- {id: guardrail.block_code_execution.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/block_code_execution", rationale: "Code-injection prevention"}
|
||||
- {id: guardrail.tool_permission.pre_call.allows, module: guardrail, tier: P1, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "guardrail_hooks/tool_permission.py", rationale: "Grant allowed tools"}
|
||||
- {id: guardrail.tool_permission.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/tool_permission.py", rationale: "Block unauthorized tools"}
|
||||
- {id: guardrail.tool_permission.pre_call.allows, module: guardrail, tier: P1, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions, messages, responses], source: "guardrail_hooks/tool_permission.py", rationale: "Grant allowed tools"}
|
||||
- {id: guardrail.tool_permission.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions, messages, responses], source: "guardrail_hooks/tool_permission.py", rationale: "Block unauthorized tools"}
|
||||
- {id: guardrail.microsoft_purview.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/microsoft_purview/purview_dlp.py", rationale: "DLP sensitive-data disclosure"}
|
||||
- {id: guardrail.headroom.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/headroom/headroom.py", rationale: "Anomaly detection threshold"}
|
||||
- {id: guardrail.generic_guardrail_api.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py", rationale: "Vendor-agnostic custom API"}
|
||||
|
|
|
|||
|
|
@ -19,10 +19,9 @@ verdict comes from the proxy itself.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Final
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import StreamingResponse, UnknownApiError
|
||||
from e2e_metadata import Capability, Domain, Mode, Provider, Subject, meta
|
||||
|
|
@ -31,14 +30,29 @@ from guardrails_client import (
|
|||
ToolPermissionParamsBody,
|
||||
ToolPermissionRuleBody,
|
||||
poll_until_blocked,
|
||||
poll_until_blocked_stream,
|
||||
poll_until_guardrail_applied,
|
||||
)
|
||||
from lifecycle import ResourceManager
|
||||
from models import ChatResponse, ChatTool, ChatToolFunction
|
||||
from models import (
|
||||
AnthropicCustomTool,
|
||||
AnthropicMessagesBody,
|
||||
AnthropicMessagesResponse,
|
||||
AnthropicToolChoice,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ChatTool,
|
||||
ChatToolFunction,
|
||||
JsonSchemaProperty,
|
||||
ResponsesFunctionTool,
|
||||
ResponsesToolBody,
|
||||
ResponsesToolOutput,
|
||||
ToolInputSchema,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
pytestmark: Final = pytest.mark.e2e
|
||||
|
||||
MODEL = "gemini-2.5-flash"
|
||||
MODEL: Final = "gemini-2.5-flash"
|
||||
|
||||
#: The one tool the guardrail permits, and one it does not. Both are declared by
|
||||
#: the caller in the request body; the guardrail reads them there.
|
||||
|
|
@ -63,11 +77,15 @@ DENIED_TOOL: Final = ChatTool(
|
|||
|
||||
TOOL_PROMPT: Final = "What is the weather in Paris right now?"
|
||||
|
||||
type FlatToolEndpoint = Literal["messages", "responses"]
|
||||
ALLOWED_TOOL_SCHEMA: Final = ToolInputSchema(properties={"city": JsonSchemaProperty(type="string")}, required=["city"])
|
||||
DENIED_TOOL_SCHEMA: Final = ToolInputSchema()
|
||||
|
||||
|
||||
def _register_tool_permission(client: GuardrailsClient, resources: ResourceManager, *, name: str) -> None:
|
||||
"""Allow-list exactly one tool: everything else falls to `default_action=deny`
|
||||
and, with `on_disallowed_action=block`, is rejected outright."""
|
||||
guardrail_id = client.register(
|
||||
guardrail_id: Final = client.register(
|
||||
name,
|
||||
ToolPermissionParamsBody(
|
||||
mode="pre_call",
|
||||
|
|
@ -100,6 +118,57 @@ def _tool_call_names(response: ChatResponse) -> tuple[str, ...]:
|
|||
)
|
||||
|
||||
|
||||
def _send_flat_tool(
|
||||
client: GuardrailsClient,
|
||||
key: str,
|
||||
endpoint: FlatToolEndpoint,
|
||||
*,
|
||||
guardrail: str,
|
||||
tool: ChatTool,
|
||||
schema: ToolInputSchema,
|
||||
require_tool: bool = False,
|
||||
) -> StreamingResponse:
|
||||
name: Final = tool.function.name
|
||||
match endpoint:
|
||||
case "messages":
|
||||
return client.proxy.transport.send(
|
||||
"/v1/messages",
|
||||
headers=client.proxy.transport.bearer(key),
|
||||
json=AnthropicMessagesBody(
|
||||
model=MODEL,
|
||||
messages=[ChatMessage(role="user", content=TOOL_PROMPT)],
|
||||
max_tokens=128,
|
||||
tools=[
|
||||
AnthropicCustomTool(name=name, description=tool.function.description or "", input_schema=schema)
|
||||
],
|
||||
tool_choice=AnthropicToolChoice(type="any") if require_tool else None,
|
||||
guardrails=[guardrail],
|
||||
),
|
||||
)
|
||||
case "responses":
|
||||
return client.proxy.transport.send(
|
||||
"/v1/responses",
|
||||
headers=client.proxy.transport.bearer(key),
|
||||
json=ResponsesToolBody(
|
||||
model=MODEL,
|
||||
input=TOOL_PROMPT,
|
||||
tools=[ResponsesFunctionTool(name=name, description=tool.function.description, parameters=schema)],
|
||||
tool_choice="required" if require_tool else None,
|
||||
guardrails=[guardrail],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _flat_tool_call_names(endpoint: FlatToolEndpoint, body: str) -> tuple[str, ...]:
|
||||
match endpoint:
|
||||
case "messages":
|
||||
content: Final = AnthropicMessagesResponse.model_validate_json(body).content or []
|
||||
return tuple(block.name for block in content if block.type == "tool_use" and block.name)
|
||||
case "responses":
|
||||
output: Final = ResponsesToolOutput.model_validate_json(body).output
|
||||
return tuple(item.name for item in output if item.type == "function_call" and item.name)
|
||||
|
||||
|
||||
class TestToolPermissionPreCall:
|
||||
@pytest.mark.covers("guardrail.tool_permission.pre_call.blocks", exercised_on=["chat_completions"])
|
||||
@meta(
|
||||
|
|
@ -118,10 +187,10 @@ class TestToolPermissionPreCall:
|
|||
rejected with a 400 that names the denied tool. An unauthorized tool that
|
||||
merely reaches the model is the whole failure mode this guardrail exists
|
||||
to prevent, so a 200 here is a hard failure."""
|
||||
name = f"e2e-toolperm-block-{unique_marker()}"
|
||||
name: Final = f"e2e-toolperm-block-{unique_marker()}"
|
||||
_register_tool_permission(client, resources, name=name)
|
||||
|
||||
result = poll_until_blocked(
|
||||
result: Final = poll_until_blocked(
|
||||
lambda: client.chat(
|
||||
scoped_key,
|
||||
MODEL,
|
||||
|
|
@ -161,10 +230,10 @@ class TestToolPermissionPreCall:
|
|||
and the model calls it. Without the header check a guardrail that never
|
||||
attached would pass this test for the wrong reason, so the 200 alone is
|
||||
not the contract."""
|
||||
name = f"e2e-toolperm-allow-{unique_marker()}"
|
||||
name: Final = f"e2e-toolperm-allow-{unique_marker()}"
|
||||
_register_tool_permission(client, resources, name=name)
|
||||
|
||||
outcome = poll_until_guardrail_applied(
|
||||
outcome: Final = poll_until_guardrail_applied(
|
||||
lambda: client.chat_raw(
|
||||
scoped_key,
|
||||
MODEL,
|
||||
|
|
@ -178,13 +247,117 @@ class TestToolPermissionPreCall:
|
|||
)
|
||||
|
||||
assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}"
|
||||
applied = _applied_guardrails(outcome)
|
||||
applied: Final = _applied_guardrails(outcome)
|
||||
assert name in applied, (
|
||||
"the allowed call must carry x-litellm-applied-guardrails naming the guardrail; "
|
||||
f"without it the 200 only proves the guardrail never ran. Got {applied!r}"
|
||||
)
|
||||
|
||||
called = _tool_call_names(ChatResponse.model_validate_json(outcome.body))
|
||||
called: Final = _tool_call_names(ChatResponse.model_validate_json(outcome.body))
|
||||
assert called == (ALLOWED_TOOL.function.name,), (
|
||||
f"the served call must carry one tool call for the permitted tool, got {called!r}: {outcome.body[:400]}"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"endpoint",
|
||||
[
|
||||
pytest.param(
|
||||
"messages",
|
||||
marks=pytest.mark.covers("guardrail.tool_permission.pre_call.blocks", exercised_on=["messages"]),
|
||||
),
|
||||
pytest.param(
|
||||
"responses",
|
||||
marks=pytest.mark.covers("guardrail.tool_permission.pre_call.blocks", exercised_on=["responses"]),
|
||||
),
|
||||
],
|
||||
)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
capabilities=(Capability.FUNCTION_CALLING,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_pre_call_blocks_flat_tool_outside_the_allow_list(
|
||||
self,
|
||||
client: GuardrailsClient,
|
||||
resources: ResourceManager,
|
||||
scoped_key: str,
|
||||
endpoint: FlatToolEndpoint,
|
||||
) -> None:
|
||||
name: Final = f"e2e-toolperm-block-{endpoint}-{unique_marker()}"
|
||||
_register_tool_permission(client, resources, name=name)
|
||||
|
||||
outcome: Final = poll_until_blocked_stream(
|
||||
lambda: _send_flat_tool(
|
||||
client, scoped_key, endpoint, guardrail=name, tool=DENIED_TOOL, schema=DENIED_TOOL_SCHEMA
|
||||
)
|
||||
)
|
||||
|
||||
assert outcome.status_code == 400, (
|
||||
f"tool_permission let a tool outside the allow-list through /{endpoint}; "
|
||||
f"got {outcome.status_code}: {outcome.body[:400]}"
|
||||
)
|
||||
assert DENIED_TOOL.function.name in outcome.body, (
|
||||
f"the block must name the denied tool so the caller can fix the request; got: {outcome.body[:400]}"
|
||||
)
|
||||
assert "Violated guardrail policy" in outcome.body, (
|
||||
f"the block must be a guardrail verdict, got: {outcome.body[:400]}"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"endpoint",
|
||||
[
|
||||
pytest.param(
|
||||
"messages",
|
||||
marks=pytest.mark.covers("guardrail.tool_permission.pre_call.allows", exercised_on=["messages"]),
|
||||
),
|
||||
pytest.param(
|
||||
"responses",
|
||||
marks=pytest.mark.covers("guardrail.tool_permission.pre_call.allows", exercised_on=["responses"]),
|
||||
),
|
||||
],
|
||||
)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
providers=(Provider.GEMINI,),
|
||||
models=(MODEL,),
|
||||
capabilities=(Capability.FUNCTION_CALLING,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_pre_call_allows_permitted_flat_tool(
|
||||
self,
|
||||
client: GuardrailsClient,
|
||||
resources: ResourceManager,
|
||||
scoped_key: str,
|
||||
endpoint: FlatToolEndpoint,
|
||||
) -> None:
|
||||
name: Final = f"e2e-toolperm-allow-{endpoint}-{unique_marker()}"
|
||||
_register_tool_permission(client, resources, name=name)
|
||||
|
||||
outcome: Final = poll_until_guardrail_applied(
|
||||
lambda: _send_flat_tool(
|
||||
client,
|
||||
scoped_key,
|
||||
endpoint,
|
||||
guardrail=name,
|
||||
tool=ALLOWED_TOOL,
|
||||
schema=ALLOWED_TOOL_SCHEMA,
|
||||
require_tool=True,
|
||||
),
|
||||
name,
|
||||
)
|
||||
|
||||
assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}"
|
||||
applied: Final = _applied_guardrails(outcome)
|
||||
assert name in applied, (
|
||||
f"the allowed call must carry x-litellm-applied-guardrails naming the guardrail; got {applied!r}"
|
||||
)
|
||||
called: Final = _flat_tool_call_names(endpoint, outcome.body)
|
||||
assert called == (ALLOWED_TOOL.function.name,), (
|
||||
f"the served call must carry one tool call for the permitted tool, got {called!r}: {outcome.body[:400]}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -653,6 +653,35 @@ class ResponsesStreamBody(BaseModel):
|
|||
cache: dict[str, bool] | None = {"no-cache": True}
|
||||
|
||||
|
||||
class ResponsesFunctionTool(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
type: Literal["function"] = "function"
|
||||
name: str
|
||||
description: str | None = None
|
||||
parameters: ToolInputSchema
|
||||
|
||||
|
||||
class ResponsesToolBody(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
model: str
|
||||
input: str
|
||||
tools: list[ResponsesFunctionTool]
|
||||
tool_choice: Literal["required"] | None = None
|
||||
guardrails: list[str]
|
||||
max_output_tokens: int = 128
|
||||
|
||||
|
||||
class ResponsesOutputItem(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
type: str
|
||||
name: str | None = None
|
||||
|
||||
|
||||
class ResponsesToolOutput(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
output: list[ResponsesOutputItem]
|
||||
|
||||
|
||||
class CountTokensBody(BaseModel):
|
||||
"""POST /v1/messages/count_tokens body: the /v1/messages shape minus
|
||||
max_tokens (the endpoint only counts the prompt)."""
|
||||
|
|
|
|||
|
|
@ -1785,6 +1785,289 @@ def test_vertex_ai_map_tools():
|
|||
assert tools == new_tools
|
||||
|
||||
|
||||
def test_vertex_ai_map_tools_declares_a_namespace_by_its_members() -> None:
|
||||
namespace = {
|
||||
"type": "namespace",
|
||||
"name": "github",
|
||||
"tools": [
|
||||
{"type": "function", "name": "create_issue", "parameters": {"type": "object", "properties": {}}},
|
||||
{"type": "function", "function": {"name": "close_issue"}},
|
||||
],
|
||||
}
|
||||
|
||||
tools = VertexGeminiConfig().map_function(value=[namespace], optional_params={})
|
||||
|
||||
assert [declaration["name"] for declaration in tools[0]["function_declarations"]] == [
|
||||
"create_issue",
|
||||
"close_issue",
|
||||
]
|
||||
|
||||
|
||||
def test_vertex_ai_map_tools_declares_a_namespace_member_like_the_same_chat_tool() -> None:
|
||||
parameters = {
|
||||
"type": "object",
|
||||
"properties": {"repo": {"$ref": "#/$defs/Repo"}, "label": {"type": ["string", "null"]}},
|
||||
"$defs": {"Repo": {"type": "string"}},
|
||||
}
|
||||
function = {"name": "create_issue", "description": "Open an issue", "parameters": parameters}
|
||||
namespace = {"type": "namespace", "name": "github", "tools": [{"type": "function", **function}]}
|
||||
|
||||
from_namespace = VertexGeminiConfig().map_function(value=[namespace], optional_params={})
|
||||
from_chat_tool = VertexGeminiConfig().map_function(
|
||||
value=[{"type": "function", "function": function}], optional_params={}
|
||||
)
|
||||
|
||||
assert from_namespace == from_chat_tool
|
||||
|
||||
|
||||
_GET_WEATHER = {"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}}
|
||||
|
||||
|
||||
_WEATHER_NAMESPACE = {"type": "namespace", "name": "weather", "tools": [_GET_WEATHER]}
|
||||
_TOP_LEVEL_GET_WEATHER = {"type": "function", "function": {k: v for k, v in _GET_WEATHER.items() if k != "type"}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tools", "declared"),
|
||||
[
|
||||
(
|
||||
[{"type": "namespace", "name": "forecast", "tools": [_GET_WEATHER]}, _WEATHER_NAMESPACE],
|
||||
["forecast__get_weather", "weather__get_weather"],
|
||||
),
|
||||
([_TOP_LEVEL_GET_WEATHER, _WEATHER_NAMESPACE], ["get_weather", "weather__get_weather"]),
|
||||
([_WEATHER_NAMESPACE, _TOP_LEVEL_GET_WEATHER], ["weather__get_weather", "get_weather"]),
|
||||
],
|
||||
ids=["namespaces_share_a_member", "top_level_first", "namespace_first"],
|
||||
)
|
||||
def test_vertex_ai_map_tools_declares_qualified_member_names_when_another_tool_declares_a_member_name(
|
||||
tools: list[dict], declared: list[str]
|
||||
) -> None:
|
||||
mapped = VertexGeminiConfig().map_function(value=tools, optional_params={})
|
||||
|
||||
assert [declaration["name"] for declaration in mapped[0]["function_declarations"]] == declared
|
||||
|
||||
|
||||
def test_vertex_ai_map_tools_declares_a_member_a_namespace_lists_twice_once_from_its_first_entry() -> None:
|
||||
first = {**_GET_WEATHER, "description": "Current weather"}
|
||||
repeat = {**_GET_WEATHER, "description": "Forecast"}
|
||||
namespace = {"type": "namespace", "name": "weather", "tools": [first, repeat]}
|
||||
|
||||
mapped = VertexGeminiConfig().map_function(value=[namespace], optional_params={})
|
||||
|
||||
declarations = mapped[0]["function_declarations"]
|
||||
assert [(declaration["name"], declaration["description"]) for declaration in declarations] == [
|
||||
("get_weather", "Current weather")
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tools", "declared"),
|
||||
[
|
||||
(
|
||||
[
|
||||
_WEATHER_NAMESPACE,
|
||||
_TOP_LEVEL_GET_WEATHER,
|
||||
{"type": "function", "function": {"name": "weather__get_weather", "description": "Top level"}},
|
||||
],
|
||||
[("get_weather", None), ("weather__get_weather", "Top level")],
|
||||
),
|
||||
(
|
||||
[
|
||||
{"type": "namespace", "name": "mcp__ops", "tools": [{"type": "function", "name": "exec"}]},
|
||||
{"type": "namespace", "name": "mcp", "tools": [{"type": "function", "name": "ops__exec"}]},
|
||||
{"type": "function", "function": {"name": "exec"}},
|
||||
{"type": "function", "function": {"name": "ops__exec"}},
|
||||
],
|
||||
[("mcp__ops__exec", None), ("exec", None), ("ops__exec", None)],
|
||||
),
|
||||
(
|
||||
[_TOP_LEVEL_GET_WEATHER, _TOP_LEVEL_GET_WEATHER, _WEATHER_NAMESPACE],
|
||||
[("get_weather", None), ("get_weather", None), ("weather__get_weather", None)],
|
||||
),
|
||||
],
|
||||
ids=["top_level_function_has_the_qualified_name", "namespaces_build_one_qualified_name", "client_repeats_a_function"],
|
||||
)
|
||||
def test_vertex_ai_map_tools_never_repeats_a_generated_name(
|
||||
tools: list[dict], declared: list[tuple[str, str | None]]
|
||||
) -> None:
|
||||
mapped = VertexGeminiConfig().map_function(value=tools, optional_params={})
|
||||
|
||||
declarations = mapped[0]["function_declarations"]
|
||||
assert [(declaration["name"], declaration.get("description")) for declaration in declarations] == declared
|
||||
|
||||
|
||||
def test_vertex_ai_map_openai_params_qualifies_members_by_the_namespace_top_level_name() -> None:
|
||||
namespace = {
|
||||
"type": "namespace",
|
||||
"name": "safe",
|
||||
"function": {"name": "lookup"},
|
||||
"tools": [{"type": "function", "name": "lookup"}],
|
||||
}
|
||||
|
||||
mapped = VertexGeminiConfig().map_openai_params(
|
||||
non_default_params={
|
||||
"tools": [namespace],
|
||||
"tool_choice": {"type": "function", "function": {"name": "lookup"}},
|
||||
},
|
||||
optional_params={},
|
||||
model="gemini-3.8-flash",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
declared = [declaration["name"] for tool in mapped["tools"] for declaration in tool["function_declarations"]]
|
||||
assert declared == ["safe__lookup"]
|
||||
assert mapped["tool_choice"]["functionCallingConfig"]["allowed_function_names"] == ["safe__lookup"]
|
||||
|
||||
|
||||
def test_vertex_ai_map_openai_params_declares_qualified_member_names_when_functions_repeat_a_member() -> None:
|
||||
mapped = VertexGeminiConfig().map_openai_params(
|
||||
non_default_params={
|
||||
"tools": [{"type": "namespace", "name": "mcp__weather", "tools": [{"type": "function", "name": "lookup"}]}],
|
||||
"functions": [{"name": "lookup", "parameters": {"type": "object", "properties": {}}}],
|
||||
},
|
||||
optional_params={},
|
||||
model="gemini-3.8-flash",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
declared = [declaration["name"] for tool in mapped["tools"] for declaration in tool["function_declarations"]]
|
||||
assert declared == ["mcp__weather__lookup", "lookup"]
|
||||
|
||||
|
||||
def test_vertex_ai_map_tools_keeps_repeated_chat_tools_when_no_namespace_is_sent() -> None:
|
||||
function = {k: v for k, v in _GET_WEATHER.items() if k != "type"}
|
||||
|
||||
mapped = VertexGeminiConfig().map_function(
|
||||
value=[{"type": "function", "function": function}, {"type": "function", "function": function}],
|
||||
optional_params={},
|
||||
)
|
||||
|
||||
assert [declaration["name"] for declaration in mapped[0]["function_declarations"]] == ["get_weather", "get_weather"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("forced_name", "other_tools", "allowed_function_names"),
|
||||
[
|
||||
("github", [], ["create_issue", "close_issue"]),
|
||||
("close_issue", [], ["close_issue"]),
|
||||
("github", [{"type": "function", "function": {"name": "github"}}], ["github"]),
|
||||
(
|
||||
"github",
|
||||
[{"type": "function", "function": {"name": "close_issue"}}],
|
||||
["github__create_issue", "github__close_issue"],
|
||||
),
|
||||
(
|
||||
"close_issue",
|
||||
[{"type": "namespace", "name": "gitlab", "tools": [{"type": "function", "name": "close_issue"}]}],
|
||||
["github__close_issue", "gitlab__close_issue"],
|
||||
),
|
||||
(
|
||||
"github",
|
||||
[
|
||||
{"type": "function", "function": {"name": "close_issue"}},
|
||||
{"type": "function", "function": {"name": "github__close_issue"}},
|
||||
],
|
||||
["github__create_issue"],
|
||||
),
|
||||
(
|
||||
"close_issue",
|
||||
[
|
||||
{"type": "namespace", "name": "gitlab", "tools": [{"type": "function", "name": "close_issue"}]},
|
||||
{"type": "function", "function": {"name": "github__close_issue"}},
|
||||
],
|
||||
["gitlab__close_issue"],
|
||||
),
|
||||
(
|
||||
"mcp",
|
||||
[
|
||||
{"type": "namespace", "name": "mcp__ops", "tools": [{"type": "function", "name": "exec"}]},
|
||||
{
|
||||
"type": "namespace",
|
||||
"name": "mcp",
|
||||
"tools": [{"type": "function", "name": "ops__exec"}, {"type": "function", "name": "status"}],
|
||||
},
|
||||
{"type": "function", "function": {"name": "exec"}},
|
||||
{"type": "function", "function": {"name": "ops__exec"}},
|
||||
],
|
||||
["mcp__status"],
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"names_the_namespace",
|
||||
"names_a_member",
|
||||
"names_a_function_sharing_the_namespace_name",
|
||||
"names_a_namespace_whose_member_a_function_repeats",
|
||||
"names_a_member_two_namespaces_share",
|
||||
"names_a_namespace_whose_generated_name_a_function_holds",
|
||||
"names_a_member_whose_generated_name_a_function_holds",
|
||||
"names_a_namespace_whose_generated_name_another_namespace_holds",
|
||||
],
|
||||
)
|
||||
def test_vertex_ai_tool_choice_forces_only_functions_gemini_is_sent(
|
||||
forced_name: str, other_tools: list[dict[str, object]], allowed_function_names: list[str]
|
||||
) -> None:
|
||||
namespace = {
|
||||
"type": "namespace",
|
||||
"name": "github",
|
||||
"tools": [{"type": "function", "name": "create_issue"}, {"type": "function", "name": "close_issue"}],
|
||||
}
|
||||
|
||||
mapped = VertexGeminiConfig().map_openai_params(
|
||||
non_default_params={
|
||||
"tools": [namespace, *other_tools],
|
||||
"tool_choice": {"type": "function", "function": {"name": forced_name}},
|
||||
},
|
||||
optional_params={},
|
||||
model="gemini-3.8-flash",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
declared = [declaration["name"] for declaration in mapped["tools"][0]["function_declarations"]]
|
||||
assert mapped["tool_choice"]["functionCallingConfig"]["allowed_function_names"] == allowed_function_names
|
||||
assert set(allowed_function_names) <= set(declared)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("forced_name", "tools"),
|
||||
[
|
||||
(
|
||||
"weather",
|
||||
[
|
||||
{"type": "namespace", "name": "weather", "tools": [{"type": "function", "name": "lookup"}]},
|
||||
{"type": "function", "function": {"name": "lookup"}},
|
||||
{"type": "function", "function": {"name": "weather__lookup"}},
|
||||
],
|
||||
),
|
||||
(
|
||||
"ops__exec",
|
||||
[
|
||||
{"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]},
|
||||
{"type": "namespace", "name": "ops__exec", "tools": [{"type": "function", "name": "status"}]},
|
||||
{"type": "function", "function": {"name": "exec"}},
|
||||
{"type": "function", "function": {"name": "status"}},
|
||||
{"type": "function", "function": {"name": "ops__exec__status"}},
|
||||
],
|
||||
),
|
||||
],
|
||||
ids=["its_members_lose_to_functions", "another_namespace_generates_its_name"],
|
||||
)
|
||||
def test_vertex_ai_tool_choice_naming_a_namespace_gemini_gets_no_member_of(
|
||||
forced_name: str, tools: list[dict[str, object]]
|
||||
) -> None:
|
||||
params = {"tools": tools, "tool_choice": {"type": "function", "function": {"name": forced_name}}}
|
||||
|
||||
with pytest.raises(litellm.UnsupportedParamsError, match=f"tool_choice '{forced_name}'"):
|
||||
VertexGeminiConfig().map_openai_params(
|
||||
non_default_params=params, optional_params={}, model="gemini-3.8-flash", drop_params=False
|
||||
)
|
||||
dropped = VertexGeminiConfig().map_openai_params(
|
||||
non_default_params=params, optional_params={}, model="gemini-3.8-flash", drop_params=True
|
||||
)
|
||||
|
||||
assert "tool_choice" not in dropped
|
||||
assert dropped["tools"][0]["function_declarations"]
|
||||
|
||||
|
||||
def test_vertex_ai_map_tool_with_anyof():
|
||||
"""
|
||||
Related issue: https://github.com/BerriAI/litellm/issues/11164
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -2556,6 +2556,66 @@ class TestResponsesWSFirstFrameModelAuth:
|
|||
"response": {"input": "hi [guarded]", "tools": tools[:1]},
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("nested", [True, False])
|
||||
async def test_frame_guard_forwards_the_tool_choice_a_guardrail_rewrote(self, nested: bool):
|
||||
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
|
||||
|
||||
async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only):
|
||||
return {**data, "tools": [], "tool_choice": "none"}
|
||||
|
||||
proxy_logging_obj = MagicMock(during_call_hook=AsyncMock())
|
||||
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook)
|
||||
guard = _responses_ws_frame_guard(
|
||||
{"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock()
|
||||
)
|
||||
body = {
|
||||
"input": "hi",
|
||||
"tools": [{"type": "function", "name": "a"}],
|
||||
"tool_choice": {"type": "function", "name": "a"},
|
||||
}
|
||||
frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body}
|
||||
|
||||
guarded = json.loads((await guard(json.dumps(frame))).frame)
|
||||
|
||||
sent = guarded["response"] if nested else guarded
|
||||
assert (sent["tools"], sent["tool_choice"]) == ([], "none")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("nested", [True, False])
|
||||
async def test_frame_guard_forwards_the_extra_body_tools_a_rewrite_guardrail_stripped(self, nested: bool):
|
||||
from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ToolPermissionGuardrail
|
||||
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
guardrail = ToolPermissionGuardrail(
|
||||
guardrail_name="ws-rewrite",
|
||||
rules=[{"id": "deny_shell", "tool_name": r"^shell$", "decision": "deny"}],
|
||||
default_action="allow",
|
||||
on_disallowed_action="rewrite",
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only):
|
||||
return await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, cache=MagicMock(), data=data, call_type=call_type
|
||||
)
|
||||
|
||||
proxy_logging_obj = MagicMock(during_call_hook=AsyncMock())
|
||||
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook)
|
||||
guard = _responses_ws_frame_guard(
|
||||
{"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock()
|
||||
)
|
||||
weather = {"type": "function", "name": "get_weather"}
|
||||
body = {"input": "hi", "extra_body": {"tools": [weather, {"type": "function", "name": "shell"}]}}
|
||||
frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body}
|
||||
|
||||
guarded = json.loads((await guard(json.dumps(frame))).frame)
|
||||
|
||||
sent = guarded["response"] if nested else guarded
|
||||
assert sent["extra_body"] == {"tools": [weather]}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_frame_guard_strips_the_proxy_owned_metadata_keys_a_client_sends(self):
|
||||
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
|
||||
|
|
@ -2991,7 +3051,7 @@ class TestResponsesWSFirstFrameModelAuth:
|
|||
assert dict(guarded.pii_tokens) == {"<EMAIL_ADDRESS_1>": "alice@example.com"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("field", ["input", "instructions", "tools"])
|
||||
@pytest.mark.parametrize("field", ["input", "instructions", "tools", "tool_choice", "extra_body"])
|
||||
async def test_frame_guard_rejects_and_bills_a_guarded_field_beside_a_nested_response(self, field: str):
|
||||
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
|
||||
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
|
||||
|
|
|
|||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -48150,7 +48150,7 @@ export interface components {
|
|||
id: string;
|
||||
/**
|
||||
* Tool Name
|
||||
* @description Regex pattern applied to the tool's function name
|
||||
* @description Regex pattern applied to the tool's name, or to its type when the tool has no name
|
||||
*/
|
||||
tool_name?: string | null;
|
||||
/**
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue