fix(memory): keep tools selectable with structured output

This commit is contained in:
moe-berri 2026-09-12 16:15:01 -07:00
parent 1d6d6c736a
commit de6178a3c5
2 changed files with 41 additions and 0 deletions

View file

@ -85,6 +85,7 @@ def inject_server_tools(
data: Mapping[str, object], route: ServerToolRoute, functions: Sequence[Mapping[str, object]], instructions: str
) -> Mapping[str, object]:
client_tools: Final = _items(data.get("tools"))
tool_choice: Final = data.get("tool_choice")
names: Final = frozenset(str(function["name"]) for function in functions)
if any(_tool_name(tool) in names for tool in client_tools):
raise ValueError("A client tool conflicts with a gateway memory tool name")
@ -109,6 +110,15 @@ def inject_server_tools(
return append_server_instructions(
{ # mutable-ok: Native provider JSON containers.
**data,
"tool_choice": tool_choice
if tool_choice is not None
else (
{ # mutable-ok: Native provider tool-choice JSON.
"type": "auto"
}
if route == "anthropic_messages"
else "auto"
),
"tools": [ # mutable-ok: Native provider JSON containers.
*client_tools,
*tools,

View file

@ -11,6 +11,37 @@ from litellm.litellm_core_utils.prompt_templates.server_tools import (
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.memory.policy import MemoryIdentity
from litellm.utils import get_optional_params
@pytest.mark.parametrize("choice", [None, "auto", "none"])
def test_structured_output_keeps_memory_tools_selectable(choice: str | None) -> None:
original: Final = {
"response_format": {
"type": "json_schema",
"json_schema": {
"name": "port",
"schema": {"type": "object", "properties": {"port": {"type": "integer"}}},
},
},
**({"tool_choice": choice} if choice is not None else {}),
}
prepared: Final = inject_server_tools(
original,
"acompletion",
({"name": "memory_search", "description": "Search", "parameters": {"type": "object"}},),
"Search memory before answering",
)
provider: Final = get_optional_params(
model="claude-sonnet-5",
custom_llm_provider="vertex_ai",
response_format=prepared["response_format"],
tools=prepared["tools"],
tool_choice=prepared.get("tool_choice"),
)
assert provider["tool_choice"] == {"type": choice or "auto"}
assert {tool["name"] for tool in provider["tools"]} == {"memory_search", "json_tool_call"}
assert original.get("tool_choice") == choice
@pytest.mark.parametrize("route", ["acompletion", "aresponses", "anthropic_messages"])