From de6178a3c5f70428c97b4adf80fdac68fac11135 Mon Sep 17 00:00:00 2001
From: moe-berri <moe@berri.ai>
Date: Sat, 12 Sep 2026 16:15:01 -0700
Subject: [PATCH] fix(memory): keep tools selectable with structured output

---
 .../prompt_templates/server_tools.py          | 10 ++++++
 .../proxy/memory/test_memory_v2_protocols.py  | 31 +++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/litellm/litellm_core_utils/prompt_templates/server_tools.py b/litellm/litellm_core_utils/prompt_templates/server_tools.py
index f4a9344a983..47de79031a1 100644
--- a/litellm/litellm_core_utils/prompt_templates/server_tools.py
+++ b/litellm/litellm_core_utils/prompt_templates/server_tools.py
@@ -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,
diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_protocols.py b/tests/test_litellm/proxy/memory/test_memory_v2_protocols.py
index 7d3a99a6111..d570dea3f19 100644
--- a/tests/test_litellm/proxy/memory/test_memory_v2_protocols.py
+++ b/tests/test_litellm/proxy/memory/test_memory_v2_protocols.py
@@ -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"])