diff --git a/litellm/litellm_core_utils/prompt_templates/server_tool_responses.py b/litellm/litellm_core_utils/prompt_templates/server_tool_responses.py index 41d77528a48..3a451f98cd3 100644 --- a/litellm/litellm_core_utils/prompt_templates/server_tool_responses.py +++ b/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") diff --git a/litellm/proxy/memory/gateway.py b/litellm/proxy/memory/gateway.py index 7c4da315b67..271c77bdab6 100644 --- a/litellm/proxy/memory/gateway.py +++ b/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(): diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py index 53a77cb6251..22fe5e1a25b 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py +++ b/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"))