fix(memory): finish replies after preparation truncation

litellm/litellm_core_utils/prompt_templates/server_tool_responses.py

@@ -138,6 +138,14 @@ def response_has_client_tools(
     )
 
 
+def response_is_truncated(response: Mapping[str, object]) -> bool:
+    return (
+        response.get("status") == "incomplete"
+        or response.get("stop_reason") == "max_tokens"
+        or any(choice.get("finish_reason") == "length" for choice in object_items(response.get("choices")))
+    )
+
+
 def executable_server_calls(
     response: Mapping[str, object], route: ServerToolRoute, server_names: frozenset[str]
 ) -> tuple[NormalizedToolCall, ...]:
@@ -163,11 +171,7 @@ def executable_server_calls(
         else bool(choices) and choices[0].get("finish_reason") == "tool_calls"
     )
     if not completed:
-        if (
-            response.get("status") == "incomplete"
-            or response.get("stop_reason") == "max_tokens"
-            or (choices and choices[0].get("finish_reason") == "length")
-        ):
+        if response_is_truncated(response):
             return ()
         raise ValueError("The model did not complete its memory tool calls")
 

litellm/proxy/memory/gateway.py

@@ -13,9 +13,9 @@ from starlette.responses import JSONResponse, Response
 from litellm._logging import verbose_proxy_logger
 from litellm.litellm_core_utils.prompt_templates.server_tool_responses import (
     executable_server_calls,
-    object_items,
     object_value,
     response_has_client_tools,
+    response_is_truncated,
     response_messages,
 )
 from litellm.litellm_core_utils.prompt_templates.server_tool_stream import ServerToolStream, ServerToolStreamError
@@ -300,6 +300,8 @@ class GatewayMemoryLoop:
         if response is None:
             raise HTTPException(status_code=502, detail="No model response received")
         self.upstream_ids = (*self.upstream_ids, str(response["id"]))
+        if self.preparing_output and response_is_truncated(response):
+            return True
         try:
             memory_calls: Final = executable_server_calls(response, self.route, MEMORY_TOOL_NAMES)
         except ValueError as exc:
@@ -341,14 +343,13 @@ class GatewayMemoryLoop:
                 yield chunk
             if await self.advance(round_index):
                 break
-        if self.preparing_output and not (
-            (self.last_response or {}).get("status") == "incomplete"
-            or (self.last_response or {}).get("stop_reason") == "max_tokens"
-            or any(
-                choice.get("finish_reason") == "length"
-                for choice in object_items((self.last_response or {}).get("choices"))
+        if self.preparing_output:
+            preparation_status: Final = (
+                "Memory preparation ran out of tokens. Do not claim a save or recall succeeded unless "
+                "the completed tool results confirm it. "
+                if self.last_response is not None and response_is_truncated(self.last_response)
+                else "Memory preparation is complete. "
             )
-        ):
             response_id: Final = self.stream.response_id
             self.stream = ServerToolStream(self.route, MEMORY_TOOL_NAMES, self.original)
             if self.route == "aresponses":
@@ -357,7 +358,7 @@ class GatewayMemoryLoop:
             self.data = append_server_reference(
                 restore_client_output(self.data, self.original),
                 self.route,
-                "Memory preparation is complete. Now respond to the user's request using the required output "
+                preparation_status + "Now respond to the user's request using the required output "
                 "format and any application tools provided. Do not describe the memory preparation.",
             )
             async for chunk in self._call():

tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py

@@ -1236,6 +1236,54 @@ async def test_private_tools_have_room_but_final_answer_keeps_client_token_cap(
         assert b"ok" in chunks and b"litellm_memory_read" not in chunks
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", _ROUTES)
+@pytest.mark.parametrize("streaming", (False, True))
+async def test_truncated_preparation_discards_partial_call_and_still_generates_final_answer(
+    prisma_edge: MagicMock, route: ServerToolRoute, streaming: bool
+) -> None:
+    observed = []
+    limit_field = "max_output_tokens" if route == "aresponses" else "max_tokens"
+    partial = {"id": "discard-partial-call", "name": "litellm_memory_capture", "arguments": {"key": "unfinished"}}
+
+    async def execute(inner: Request, body: dict[str, object], auth: UserAPIKeyAuth) -> Response:
+        observed.append(body)
+        reply = (
+            provider_response(route, "", (partial,), truncated=True)
+            if len(observed) == 1
+            else provider_response(route, "Memory was not saved")
+        )
+        return wire_response(reply, route, body.get("stream") is True)
+
+    loop = GatewayMemoryLoop(
+        execute,
+        request(),
+        {
+            "input": "Remember this",
+            "messages": [{"role": "user", "content": "Remember this"}],
+            limit_field: 40,
+            "stream": streaming,
+        },
+        route,
+        store(prisma_edge),
+        UserAPIKeyAuth(),
+    )
+    chunks = b"".join([chunk async for chunk in loop.run()])
+    assert [body[limit_field] for body in observed] == [4096, 40]
+    assert "discard-partial-call" not in json.dumps(observed[-1])
+    assert "ran out of tokens" in json.dumps(observed[-1])
+    assert "Memory was not saved" in json.dumps(loop.stream.response())
+    prisma_edge.db.litellm_memorytable.create.assert_not_awaited()
+    if route == "aresponses":
+        saved = json.loads(
+            prisma_edge.db.litellm_memorycontinuation.upsert.call_args.kwargs["data"]["create"]["payload"]
+        )
+        assert "discard-partial-call" not in json.dumps(saved)
+        assert saved["response"]["status"] == "completed"
+    if streaming:
+        assert b"Memory was not saved" in chunks and b"discard-partial-call" not in chunks
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("streaming", (False, True))
 @pytest.mark.parametrize("tool_name", ("litellm_memory_capture", "read_file"))