scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(gemini): keep generated member names unique and let a top-level function own its exact name

Browse source 
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.
This commit is contained in:
mateo-berri 2026-11-03 13:21:40 -08:00
parent 07a6a21102
commit 427f2ce4e2
4 changed files with 111 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -563,7 +563,19 @@ 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 tuple(map(_as_chat_tool, members)) if isinstance(members, list) else ()
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:
@ -598,20 +610,31 @@ def _qualified_member(namespace: str, member: JsonValue) -> JsonValue:
return {**member, "function": {**function, "name": qualified_member_name(namespace, name)}}
def _as_declared(tool: JsonValue, declarable: Counter[str | None]) -> tuple[JsonValue, ...]:
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 = declared_function_name(tool)
if not _is_namespace(tool) or namespace is None:
return (tool,)
return tuple(_qualified_member(namespace, member) for member in _namespace_members_or_tool(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
declarable: Final = Counter(_declarable_names_in((*tools, *_listed(outside_declarations))))
return list(chain.from_iterable(_as_declared(tool, declarable) for tool in 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 _qualified_names_of(member_name: str, tool: JsonValue, declarable: Counter[str | None]) -> tuple[str, ...]:
@ -629,7 +652,7 @@ def _forced_function_names(name: str, tools: JsonValue) -> tuple[str, ...]:
declarable: Final = Counter(_declarable_names_in(candidates))
named: Final = next((tool for tool in candidates if _is_namespace(tool) and _member_name_of(tool) == name), None)
if named is not None:
declared: Final = map(declared_function_name, _as_declared(named, declarable))
declared: Final = map(declared_function_name, _as_declared(named, declarable, frozenset()))
return tuple(member for member in declared if member is not None) or (name,)
qualified: Final = chain.from_iterable(_qualified_names_of(name, tool, declarable) for tool in candidates)
return tuple(qualified) or (name,)

2
litellm/proxy/guardrails/guardrail_hooks/tool_permission.py
View file

@ -1015,7 +1015,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
returned_members: Final = _grouped_values(
chain(
((bare, (name, bare)) for name, bare in member_names if bare not in top_level_names),
((name, (name, bare)) for name, bare in member_names),
((name, (name, bare)) for name, bare in member_names if name not in top_level_names),
)
)
return MappingProxyType(

55
tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -1836,12 +1836,8 @@ _TOP_LEVEL_GET_WEATHER = {"type": "function", "function": {k: v for k, v in _GET
),
([_TOP_LEVEL_GET_WEATHER, _WEATHER_NAMESPACE], ["get_weather", "weather__get_weather"]),
([_WEATHER_NAMESPACE, _TOP_LEVEL_GET_WEATHER], ["weather__get_weather", "get_weather"]),
(
[{"type": "namespace", "name": "delete_db", "tools": [_GET_WEATHER, _GET_WEATHER]}],
["delete_db__get_weather", "delete_db__get_weather"],
),
],
ids=["namespaces_share_a_member", "top_level_first", "namespace_first", "namespace_repeats_a_member"],
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]
@ -1851,6 +1847,55 @@ def test_vertex_ai_map_tools_declares_qualified_member_names_when_another_tool_d
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_declares_qualified_member_names_when_functions_repeat_a_member() -> None:
mapped = VertexGeminiConfig().map_openai_params(
non_default_params={

31
tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py
View file

@ -706,6 +706,37 @@ class TestToolPermissionGuardrail:
assert [call.function.name for call in response.choices[0].message.tool_calls or ()] == kept
@pytest.mark.asyncio
async def test_post_call_judges_a_top_level_function_by_its_own_name_beside_a_member_qualified_the_same(
self,
) -> None:
guardrail = ToolPermissionGuardrail(
guardrail_name="post-exec",
rules=[{"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}],
default_action="allow",
on_disallowed_action="rewrite",
event_hook=GuardrailEventHooks.post_call,
)
namespace = {"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]}
top_level = {"type": "function", "function": {"name": "ops__exec"}}
function = Function(name="ops__exec", arguments="{}")
response = ModelResponse(
choices=[
Choices(
finish_reason="tool_calls",
message={"tool_calls": [{"id": "call_ops", "type": "function", "function": function}]},
)
]
)
await guardrail.async_post_call_success_hook(
data={"guardrails": ["post-exec"], "tools": [namespace, top_level]},
user_api_key_dict=UserAPIKeyAuth(),
response=response,
)
assert [call.function.name for call in response.choices[0].message.tool_calls or ()] == ["ops__exec"]
@pytest.mark.asyncio
async def test_async_post_call_success_hook_missing_arguments_blocks_param_rule(
self,