From 20279d9a50ac051359d791dea77865c17c0bbf9e Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Fri, 27 Jun 2025 10:44:39 -0700 Subject: [PATCH] fix: Save output in state via output_key only when the event is authored by current agent PiperOrigin-RevId: 776640671 --- src/google/adk/agents/llm_agent.py | 10 + .../agents/test_llm_agent_output_save.py | 210 ++++++++++++++++++ 2 files changed, 220 insertions(+) create mode 100644 tests/unittests/agents/test_llm_agent_output_save.py diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py index 64c3628d..74091bf2 100644 --- a/src/google/adk/agents/llm_agent.py +++ b/src/google/adk/agents/llm_agent.py @@ -431,12 +431,22 @@ class LlmAgent(BaseAgent): def __maybe_save_output_to_state(self, event: Event): """Saves the model output to state if needed.""" + # skip if the event was authored by some other agent (e.g. current agent + # transferred to another agent) + if event.author != self.name: + logger.debug( + 'Skipping output save for agent %s: event authored by %s', + self.name, + event.author, + ) + return if ( self.output_key and event.is_final_response() and event.content and event.content.parts ): + result = ''.join( [part.text if part.text else '' for part in event.content.parts] ) diff --git a/tests/unittests/agents/test_llm_agent_output_save.py b/tests/unittests/agents/test_llm_agent_output_save.py new file mode 100644 index 00000000..a13fde84 --- /dev/null +++ b/tests/unittests/agents/test_llm_agent_output_save.py @@ -0,0 +1,210 @@ +# Copyright 2025 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. + +"""Unit tests for LlmAgent output saving functionality.""" + +from unittest.mock import Mock +from unittest.mock import patch + +from google.adk.agents.llm_agent import LlmAgent +from google.adk.events.event import Event +from google.adk.events.event_actions import EventActions +from google.genai import types +from pydantic import BaseModel +import pytest + + +class MockOutputSchema(BaseModel): + message: str + confidence: float + + +def create_test_event( + author: str = "test_agent", + content_text: str = "Hello world", + is_final: bool = True, + invocation_id: str = "test_invocation", +) -> Event: + """Helper to create test events.""" + # Create mock content + parts = [types.Part.from_text(text=content_text)] if content_text else [] + content = types.Content(role="model", parts=parts) if parts else None + + # Create event + event = Event( + invocation_id=invocation_id, + author=author, + content=content, + actions=EventActions(), + ) + + # Mock is_final_response if needed + if not is_final: + event.partial = True + + return event + + +class TestLlmAgentOutputSave: + """Test suite for LlmAgent output saving functionality.""" + + def test_maybe_save_output_to_state_skips_different_author(self, caplog): + """Test that output is not saved when event author differs from agent name.""" + agent = LlmAgent(name="agent_a", output_key="result") + event = create_test_event(author="agent_b", content_text="Response from B") + + with caplog.at_level("DEBUG"): + agent._LlmAgent__maybe_save_output_to_state(event) + + # Should not add anything to state_delta + assert len(event.actions.state_delta) == 0 + + # Should log the skip + assert ( + "Skipping output save for agent agent_a: event authored by agent_b" + in caplog.text + ) + + def test_maybe_save_output_to_state_saves_same_author(self): + """Test that output is saved when event author matches agent name.""" + agent = LlmAgent(name="test_agent", output_key="result") + event = create_test_event(author="test_agent", content_text="Test response") + + agent._LlmAgent__maybe_save_output_to_state(event) + + # Should save to state_delta + assert event.actions.state_delta["result"] == "Test response" + + def test_maybe_save_output_to_state_no_output_key(self): + """Test that nothing is saved when output_key is not set.""" + agent = LlmAgent(name="test_agent") # No output_key + event = create_test_event(author="test_agent", content_text="Test response") + + agent._LlmAgent__maybe_save_output_to_state(event) + + # Should not save anything + assert len(event.actions.state_delta) == 0 + + def test_maybe_save_output_to_state_not_final_response(self): + """Test that output is not saved for non-final responses.""" + agent = LlmAgent(name="test_agent", output_key="result") + event = create_test_event( + author="test_agent", content_text="Partial response", is_final=False + ) + + agent._LlmAgent__maybe_save_output_to_state(event) + + # Should not save partial responses + assert len(event.actions.state_delta) == 0 + + def test_maybe_save_output_to_state_no_content(self): + """Test that nothing is saved when event has no content.""" + agent = LlmAgent(name="test_agent", output_key="result") + event = create_test_event(author="test_agent", content_text="") + + agent._LlmAgent__maybe_save_output_to_state(event) + + # Should not save empty content + assert len(event.actions.state_delta) == 0 + + def test_maybe_save_output_to_state_with_output_schema(self): + """Test that output is processed with schema when output_schema is set.""" + agent = LlmAgent( + name="test_agent", output_key="result", output_schema=MockOutputSchema + ) + + # Create event with JSON content + json_content = '{"message": "Hello", "confidence": 0.95}' + event = create_test_event(author="test_agent", content_text=json_content) + + agent._LlmAgent__maybe_save_output_to_state(event) + + # Should save parsed and validated output + expected_output = {"message": "Hello", "confidence": 0.95} + assert event.actions.state_delta["result"] == expected_output + + def test_maybe_save_output_to_state_multiple_parts(self): + """Test that multiple text parts are concatenated.""" + agent = LlmAgent(name="test_agent", output_key="result") + + # Create event with multiple text parts + parts = [ + types.Part.from_text(text="Hello "), + types.Part.from_text(text="world"), + types.Part.from_text(text="!"), + ] + content = types.Content(role="model", parts=parts) + + event = Event( + invocation_id="test_invocation", + author="test_agent", + content=content, + actions=EventActions(), + ) + + agent._LlmAgent__maybe_save_output_to_state(event) + + # Should concatenate all text parts + assert event.actions.state_delta["result"] == "Hello world!" + + def test_maybe_save_output_to_state_agent_transfer_scenario(self, caplog): + """Test realistic agent transfer scenario.""" + # Scenario: Agent A transfers to Agent B, Agent B produces output + # Agent A should not save Agent B's output + + agent_a = LlmAgent(name="support_agent", output_key="support_result") + agent_b_event = create_test_event( + author="billing_agent", content_text="Your bill is $100" + ) + + with caplog.at_level("DEBUG"): + agent_a._LlmAgent__maybe_save_output_to_state(agent_b_event) + + # Agent A should not save Agent B's output + assert len(agent_b_event.actions.state_delta) == 0 + assert ( + "Skipping output save for agent support_agent: event authored by" + " billing_agent" + in caplog.text + ) + + def test_maybe_save_output_to_state_case_sensitive_names(self, caplog): + """Test that agent name comparison is case-sensitive.""" + agent = LlmAgent(name="TestAgent", output_key="result") + event = create_test_event(author="testagent", content_text="Test response") + + with caplog.at_level("DEBUG"): + agent._LlmAgent__maybe_save_output_to_state(event) + + # Should not save due to case mismatch + assert len(event.actions.state_delta) == 0 + assert ( + "Skipping output save for agent TestAgent: event authored by testagent" + in caplog.text + ) + + @patch("google.adk.agents.llm_agent.logger") + def test_maybe_save_output_to_state_logging(self, mock_logger): + """Test that debug logging works correctly.""" + agent = LlmAgent(name="agent1", output_key="result") + event = create_test_event(author="agent2", content_text="Test response") + + agent._LlmAgent__maybe_save_output_to_state(event) + + # Should call logger.debug with correct parameters + mock_logger.debug.assert_called_once_with( + "Skipping output save for agent %s: event authored by %s", + "agent1", + "agent2", + )