From be103fbd2f2ef4804afd691e37de53def2b4f934 Mon Sep 17 00:00:00 2001 From: George Weale Date: Wed, 12 Aug 2026 16:10:18 -0700 Subject: [PATCH] fix: keep preloaded context turn-scoped PreloadMemoryTool appended recalled memory to system_instruction, but explicit caches fingerprint and store that system prefix, so memory picked for one query could destabilize cache identity or stick to a reusable prefix on later turns. This inserts recalled memory into the trailing user-context batch for the current request instead, which the cache manager already excludes, through one typed internal LlmRequest seam that dynamic user instructions also reuse. Co-authored-by: George Weale PiperOrigin-RevId: 963713313 --- src/google/adk/flows/llm_flows/contents.py | 34 +--- src/google/adk/models/llm_request.py | 26 +++ src/google/adk/tools/preload_memory_tool.py | 9 +- .../unittests/tools/test_load_memory_tool.py | 11 +- .../tools/test_preload_memory_tool.py | 151 ++++++++++++++++++ 5 files changed, 193 insertions(+), 38 deletions(-) create mode 100644 tests/unittests/tools/test_preload_memory_tool.py diff --git a/src/google/adk/flows/llm_flows/contents.py b/src/google/adk/flows/llm_flows/contents.py index f51211d7..cc949a9d 100644 --- a/src/google/adk/flows/llm_flows/contents.py +++ b/src/google/adk/flows/llm_flows/contents.py @@ -1306,16 +1306,6 @@ def _is_live_model_media_event_with_inline_data(event: Event) -> bool: return False -def _content_contains_function_response(content: types.Content) -> bool: - """Checks whether the content includes any function response parts.""" - if not content.parts: - return False - for part in content.parts: - if part.function_response: - return True - return False - - def _add_model_input_context_to_user_content( invocation_context: InvocationContext, llm_request: LlmRequest, @@ -1356,24 +1346,6 @@ async def _add_instructions_to_user_content( """ if not instruction_contents: return - - # Find the insertion point: before the last continuous batch of user content - # Walk backwards to find the first non-user content, then insert after it - insert_index = len(llm_request.contents) - - if llm_request.contents: - for i in range(len(llm_request.contents) - 1, -1, -1): - content = llm_request.contents[i] - if content.role != 'user': - insert_index = i + 1 - break - if _content_contains_function_response(content): - insert_index = i + 1 - break - insert_index = i - else: - # No contents remaining, just append at the end - insert_index = 0 - - # Insert all instruction contents at the proper position using efficient slicing - llm_request.contents[insert_index:insert_index] = instruction_contents + llm_request._insert_transient_user_content( # pylint: disable=protected-access + instruction_contents + ) diff --git a/src/google/adk/models/llm_request.py b/src/google/adk/models/llm_request.py index c7e0479d..e7255503 100644 --- a/src/google/adk/models/llm_request.py +++ b/src/google/adk/models/llm_request.py @@ -299,6 +299,32 @@ class LlmRequest(BaseModel): # No existing tool with function_declarations, create new one self.config.tools.append(types.Tool(function_declarations=declarations)) + def _insert_transient_user_content( + self, contents: list[types.Content] + ) -> None: + """Insert request-scoped user context at the current-turn boundary. + + Transient retrieval or dynamic instruction content belongs before the + latest ordinary user batch, but after a function response when the model + is continuing a tool-call turn. Keeping it at this boundary prevents the + request-scoped content from entering a reusable system/history prefix. + """ + if not contents: + return + + insert_index = len(self.contents) + for i in range(len(self.contents) - 1, -1, -1): + content = self.contents[i] + if content.role != "user": + insert_index = i + 1 + break + if any(part.function_response for part in content.parts or []): + insert_index = i + 1 + break + insert_index = i + + self.contents[insert_index:insert_index] = contents + def set_output_schema( self, output_schema: Optional[SchemaType] = None, diff --git a/src/google/adk/tools/preload_memory_tool.py b/src/google/adk/tools/preload_memory_tool.py index a7ae8575..a69421f0 100644 --- a/src/google/adk/tools/preload_memory_tool.py +++ b/src/google/adk/tools/preload_memory_tool.py @@ -17,6 +17,7 @@ from __future__ import annotations import logging from typing import TYPE_CHECKING +from google.genai import types from typing_extensions import override from . import _memory_entry_utils @@ -80,13 +81,17 @@ class PreloadMemoryTool(BaseTool): return full_memory_text = '\n'.join(memory_text_lines) - si = f"""The following content is from your previous conversations with the user. + memory_context = f"""The following content is from your previous conversations with the user. They may be useful for answering the user's current query. {full_memory_text} """ - llm_request._append_dynamic_instructions([si]) + llm_request._insert_transient_user_content([ # pylint: disable=protected-access + types.Content( + role='user', parts=[types.Part.from_text(text=memory_context)] + ) + ]) preload_memory_tool = PreloadMemoryTool() diff --git a/tests/unittests/tools/test_load_memory_tool.py b/tests/unittests/tools/test_load_memory_tool.py index 81d95439..7aa767a9 100644 --- a/tests/unittests/tools/test_load_memory_tool.py +++ b/tests/unittests/tools/test_load_memory_tool.py @@ -99,8 +99,8 @@ async def test_process_llm_request_appends_to_existing_system_instruction(): @pytest.mark.asyncio -async def test_preload_memory_registers_dynamic_instructions(): - """Test that PreloadMemoryTool registers memory into _dynamic_instructions.""" +async def test_preload_memory_registers_transient_user_content(): + """Test that PreloadMemoryTool registers memory as transient user content.""" tool = PreloadMemoryTool() tool_context = mock.Mock(spec=ToolContext) tool_context.user_content = types.Content( @@ -125,7 +125,8 @@ async def test_preload_memory_registers_dynamic_instructions(): tool_context=tool_context, llm_request=llm_request ) - assert len(llm_request._dynamic_instructions) == 1 - assert '' in llm_request._dynamic_instructions[0] + assert len(llm_request._dynamic_instructions) == 0 assert llm_request.config.system_instruction is None - assert len(llm_request.contents) == 0 + assert len(llm_request.contents) == 1 + assert llm_request.contents[0].role == 'user' + assert '' in llm_request.contents[0].parts[0].text diff --git a/tests/unittests/tools/test_preload_memory_tool.py b/tests/unittests/tools/test_preload_memory_tool.py new file mode 100644 index 00000000..53173e0e --- /dev/null +++ b/tests/unittests/tools/test_preload_memory_tool.py @@ -0,0 +1,151 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock + +from google.adk.memory.base_memory_service import SearchMemoryResponse +from google.adk.memory.memory_entry import MemoryEntry +from google.adk.models.gemini_context_cache_manager import GeminiContextCacheManager +from google.adk.models.llm_request import LlmRequest +from google.adk.tools.preload_memory_tool import PreloadMemoryTool +from google.genai import types +import pytest + + +def _tool_context(*memories: MemoryEntry): + tool_context = mock.Mock() + tool_context.user_content = types.UserContent('current query') + tool_context.search_memory = mock.AsyncMock( + return_value=SearchMemoryResponse(memories=list(memories)) + ) + return tool_context + + +def _memory(text: str) -> MemoryEntry: + return MemoryEntry( + content=types.UserContent(text), + author='user', + timestamp='2026-07-13T12:00:00Z', + ) + + +@pytest.mark.asyncio +async def test_preload_memory_keeps_system_prefix_stable(): + """Recalled memory goes into contents, never into the system instruction.""" + request = LlmRequest( + contents=[ + types.UserContent('historical question'), + types.ModelContent('historical answer'), + types.UserContent('current query'), + ] + ) + request.config.system_instruction = 'stable instruction' + + await PreloadMemoryTool().process_llm_request( + tool_context=_tool_context(_memory('likes tea')), + llm_request=request, + ) + + assert request.config.system_instruction == 'stable instruction' + assert [content.role for content in request.contents] == [ + 'user', + 'model', + 'user', + 'user', + ] + assert 'likes tea' in request.contents[-2].parts[0].text + assert request.contents[-1] == types.UserContent('current query') + + +@pytest.mark.asyncio +async def test_preload_memory_stays_after_function_response_boundary(): + """Recalled memory lands after a trailing function response.""" + function_response = types.Content( + role='user', + parts=[ + types.Part.from_function_response( + name='lookup', response={'result': 'done'} + ) + ], + ) + request = LlmRequest( + contents=[ + types.UserContent('current query'), + types.ModelContent( + types.Part.from_function_call(name='lookup', args={}) + ), + function_response, + ] + ) + + await PreloadMemoryTool().process_llm_request( + tool_context=_tool_context(_memory('likes tea')), + llm_request=request, + ) + + assert request.contents[-2] is function_response + assert 'likes tea' in request.contents[-1].parts[0].text + + +@pytest.mark.asyncio +async def test_preload_memory_does_not_change_cacheable_prefix_fingerprint(): + """Different recalled memories keep the same prefix fingerprint.""" + requests = [] + for memory_text in ('likes tea', 'likes coffee'): + request = LlmRequest( + model='gemini-2.5-flash', + contents=[ + types.UserContent('historical question'), + types.ModelContent('historical answer'), + types.UserContent('current query'), + ], + ) + request.config.system_instruction = 'stable instruction' + await PreloadMemoryTool().process_llm_request( + tool_context=_tool_context(_memory(memory_text)), + llm_request=request, + ) + requests.append(request) + + client = mock.Mock(vertexai=False) + client._api_client = None + manager = GeminiContextCacheManager(client) + prefix_counts = [ + manager._find_count_of_contents_to_cache(request.contents) + for request in requests + ] + fingerprints = [ + manager._generate_cache_fingerprint(request, prefix_count) + for request, prefix_count in zip(requests, prefix_counts) + ] + + assert prefix_counts == [2, 2] + assert fingerprints[0] == fingerprints[1] + + +@pytest.mark.asyncio +async def test_preload_memory_search_failure_is_noop(): + """A failing memory search leaves the request completely untouched.""" + request = LlmRequest(contents=[types.UserContent('current query')]) + request.config.system_instruction = 'stable instruction' + original = request.model_copy(deep=True) + tool_context = _tool_context() + tool_context.search_memory.side_effect = RuntimeError('unavailable') + + await PreloadMemoryTool().process_llm_request( + tool_context=tool_context, + llm_request=request, + ) + + assert request == original