diff --git a/src/google/adk/plugins/multimodal_tool_results_plugin.py b/src/google/adk/plugins/multimodal_tool_results_plugin.py index a103b52a66..dc43f68929 100644 --- a/src/google/adk/plugins/multimodal_tool_results_plugin.py +++ b/src/google/adk/plugins/multimodal_tool_results_plugin.py @@ -15,6 +15,7 @@ from __future__ import annotations from typing import Any +from typing import Literal from typing import Optional from google.genai import types @@ -27,6 +28,14 @@ from .base_plugin import BasePlugin PARTS_RETURNED_BY_TOOLS_ID = "temp:PARTS_RETURNED_BY_TOOLS_ID" +# Deliberately NOT "temp:"-prefixed: the session layer treats "temp:" state +# as invocation-scoped and strips it before persisting an event (see +# BaseSessionService._trim_temp_delta_state). retention="session" needs the +# saved parts to survive into later invocations (i.e. later conversational +# turns), so it is stored under this session-scoped key instead. +SESSION_PARTS_RETURNED_BY_TOOLS_ID = ( + "multimodal_tool_results_plugin:PARTS_RETURNED_BY_TOOLS_ID" +) class MultimodalToolResultsPlugin(BasePlugin): @@ -36,13 +45,29 @@ class MultimodalToolResultsPlugin(BasePlugin): are supported outside of computer use tool. """ - def __init__(self, name: str = "multimodal_tool_results_plugin"): + def __init__( + self, + name: str = "multimodal_tool_results_plugin", + retention: Literal["next_model_call", "session"] = "next_model_call", + ): """Initialize the multimodal tool results plugin. Args: name: The name of the plugin instance. + retention: How long tool-returned parts stay attached to model + requests. "next_model_call" (default) attaches the saved parts once + and then clears them. "session" keeps re-attaching the latest saved + parts to every subsequent model request for the rest of the + session, so follow-up turns can still reference them. """ super().__init__(name) + self._retention = retention + + def _state_key(self) -> str: + """Returns the state key parts are stored under for this retention mode.""" + if self._retention == "session": + return SESSION_PARTS_RETURNED_BY_TOOLS_ID + return PARTS_RETURNED_BY_TOOLS_ID async def after_tool_callback( self, @@ -67,11 +92,12 @@ async def after_tool_callback( return result parts = [result] if isinstance(result, types.Part) else result[:] + key = self._state_key() - if PARTS_RETURNED_BY_TOOLS_ID in tool_context.state: - tool_context.state[PARTS_RETURNED_BY_TOOLS_ID] += parts + if key in tool_context.state: + tool_context.state[key] += parts else: - tool_context.state[PARTS_RETURNED_BY_TOOLS_ID] = parts + tool_context.state[key] = parts return None @@ -83,10 +109,10 @@ async def before_model_callback( if not llm_request.contents: return None - if saved_parts := callback_context.state.get( - PARTS_RETURNED_BY_TOOLS_ID, None - ): + key = self._state_key() + if saved_parts := callback_context.state.get(key, None): llm_request.contents[-1].parts += saved_parts - callback_context.state.update({PARTS_RETURNED_BY_TOOLS_ID: []}) + if self._retention == "next_model_call": + callback_context.state.update({key: []}) return None diff --git a/tests/unittests/plugins/test_multimodal_tool_results_plugin.py b/tests/unittests/plugins/test_multimodal_tool_results_plugin.py index 1c0b6a0b5a..6386de4c9e 100644 --- a/tests/unittests/plugins/test_multimodal_tool_results_plugin.py +++ b/tests/unittests/plugins/test_multimodal_tool_results_plugin.py @@ -19,6 +19,7 @@ from google.adk.agents.base_agent import BaseAgent from google.adk.agents.callback_context import CallbackContext +from google.adk.agents.llm_agent import Agent from google.adk.models.llm_request import LlmRequest from google.adk.plugins.multimodal_tool_results_plugin import MultimodalToolResultsPlugin from google.adk.plugins.multimodal_tool_results_plugin import PARTS_RETURNED_BY_TOOLS_ID @@ -144,6 +145,53 @@ async def test_empty_contents_leaves_saved_parts_pending( assert tool_context.state[PARTS_RETURNED_BY_TOOLS_ID] == parts +@pytest.mark.asyncio +async def test_session_retention_reattaches_parts_across_turns(): + """Test that retention="session" resurfaces parts in a LATER invocation. + + Regression test for #6695. This must go through two separate + runner.run_async() calls against a real session service: the saved parts + are stored under a "temp:"-prefixed key by default, and that prefix is + stripped by BaseSessionService before an event is persisted, so a test + that only calls before_model_callback twice on the same in-memory State + object (without an intervening append_event()) cannot detect whether the + parts actually survive a real turn boundary. + """ + file_part = types.Part( + file_data=types.FileData( + file_uri="gs://bucket/document.pdf", mime_type="application/pdf" + ) + ) + + def get_document() -> types.Part: + return file_part + + mock_model = testing_utils.MockModel.create( + responses=[ + types.Part.from_function_call(name="get_document", args={}), + "Here is a summary of the document.", + "The document says X.", + ] + ) + agent = Agent(name="root_agent", model=mock_model, tools=[get_document]) + runner = testing_utils.InMemoryRunner( + agent, plugins=[MultimodalToolResultsPlugin(retention="session")] + ) + + # Turn 1: triggers the tool call. + await runner.run_async("Please fetch the document") + # Turn 2: a NEW invocation, sharing the same session as turn 1. + await runner.run_async("What does the document say?") + + assert len(mock_model.requests) == 3 + # Turn 1's first request precedes the tool call: nothing attached yet. + assert file_part not in mock_model.requests[0].contents[-1].parts + # Turn 1's second request: parts attached within the same invocation. + assert file_part in mock_model.requests[1].contents[-1].parts + # Turn 2's request is a separate invocation: parts must still be attached. + assert file_part in mock_model.requests[2].contents[-1].parts + + @pytest.mark.asyncio async def test_multiple_tools_returning_parts_are_accumulated( plugin: ToolReturningGenAiPartsPlugin,