diff --git a/src/google/adk/agents/context.py b/src/google/adk/agents/context.py new file mode 100644 index 00000000..79942681 --- /dev/null +++ b/src/google/adk/agents/context.py @@ -0,0 +1,373 @@ +# 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 __future__ import annotations + +from collections.abc import Mapping +from collections.abc import Sequence +from typing import Any +from typing import TYPE_CHECKING + +from typing_extensions import override + +from .readonly_context import ReadonlyContext + +if TYPE_CHECKING: + from google.genai import types + + from ..artifacts.base_artifact_service import ArtifactVersion + from ..auth.auth_credential import AuthCredential + from ..auth.auth_tool import AuthConfig + from ..events.event import Event + from ..events.event_actions import EventActions + from ..memory.base_memory_service import SearchMemoryResponse + from ..sessions.state import State + from ..tools.tool_confirmation import ToolConfirmation + from .invocation_context import InvocationContext + + +class Context(ReadonlyContext): + """The context within an agent run.""" + + def __init__( + self, + invocation_context: InvocationContext, + *, + event_actions: EventActions | None = None, + function_call_id: str | None = None, + tool_confirmation: ToolConfirmation | None = None, + ) -> None: + """Initializes the Context. + + Args: + invocation_context: The invocation context. + event_actions: The event actions for state and artifact deltas. + function_call_id: The function call id of the current tool call. Required + for tool-specific methods like request_credential and + request_confirmation. + tool_confirmation: The tool confirmation of the current tool call. + """ + super().__init__(invocation_context) + + from ..events.event_actions import EventActions + from ..sessions.state import State + + self._event_actions = event_actions or EventActions() + self._state = State( + value=invocation_context.session.state, + delta=self._event_actions.state_delta, + ) + self._function_call_id = function_call_id + self._tool_confirmation = tool_confirmation + + @property + def function_call_id(self) -> str | None: + """The function call id of the current tool call.""" + return self._function_call_id + + @property + def tool_confirmation(self) -> ToolConfirmation | None: + """The tool confirmation of the current tool call.""" + return self._tool_confirmation + + @property + @override + def state(self) -> State: + """The delta-aware state of the current session. + + For any state change, you can mutate this object directly, + e.g. `ctx.state['foo'] = 'bar'` + """ + return self._state + + @property + def actions(self) -> EventActions: + """The event actions for the current context.""" + return self._event_actions + + # ============================================================================ + # Artifact methods + # ============================================================================ + + async def load_artifact( + self, filename: str, version: int | None = None + ) -> types.Part | None: + """Loads an artifact attached to the current session. + + Args: + filename: The filename of the artifact. + version: The version of the artifact. If None, the latest version will be + returned. + + Returns: + The artifact. + """ + if self._invocation_context.artifact_service is None: + raise ValueError("Artifact service is not initialized.") + return await self._invocation_context.artifact_service.load_artifact( + app_name=self._invocation_context.app_name, + user_id=self._invocation_context.user_id, + session_id=self._invocation_context.session.id, + filename=filename, + version=version, + ) + + async def save_artifact( + self, + filename: str, + artifact: types.Part, + custom_metadata: dict[str, Any] | None = None, + ) -> int: + """Saves an artifact and records it as delta for the current session. + + Args: + filename: The filename of the artifact. + artifact: The artifact to save. + custom_metadata: Custom metadata to associate with the artifact. + + Returns: + The version of the artifact. + """ + if self._invocation_context.artifact_service is None: + raise ValueError("Artifact service is not initialized.") + version = await self._invocation_context.artifact_service.save_artifact( + app_name=self._invocation_context.app_name, + user_id=self._invocation_context.user_id, + session_id=self._invocation_context.session.id, + filename=filename, + artifact=artifact, + custom_metadata=custom_metadata, + ) + self._event_actions.artifact_delta[filename] = version + return version + + async def get_artifact_version( + self, filename: str, version: int | None = None + ) -> ArtifactVersion | None: + """Gets artifact version info. + + Args: + filename: The filename of the artifact. + version: The version of the artifact. If None, the latest version will be + returned. + + Returns: + The artifact version info. + """ + if self._invocation_context.artifact_service is None: + raise ValueError("Artifact service is not initialized.") + return await self._invocation_context.artifact_service.get_artifact_version( + app_name=self._invocation_context.app_name, + user_id=self._invocation_context.user_id, + session_id=self._invocation_context.session.id, + filename=filename, + version=version, + ) + + async def list_artifacts(self) -> list[str]: + """Lists the filenames of the artifacts attached to the current session.""" + if self._invocation_context.artifact_service is None: + raise ValueError("Artifact service is not initialized.") + return await self._invocation_context.artifact_service.list_artifact_keys( + app_name=self._invocation_context.app_name, + user_id=self._invocation_context.user_id, + session_id=self._invocation_context.session.id, + ) + + # ============================================================================ + # Credential methods + # ============================================================================ + + async def save_credential(self, auth_config: AuthConfig) -> None: + """Saves a credential to the credential service. + + Args: + auth_config: The authentication configuration containing the credential. + """ + if self._invocation_context.credential_service is None: + raise ValueError("Credential service is not initialized.") + await self._invocation_context.credential_service.save_credential( + auth_config, self + ) + + async def load_credential( + self, auth_config: AuthConfig + ) -> AuthCredential | None: + """Loads a credential from the credential service. + + Args: + auth_config: The authentication configuration for the credential. + + Returns: + The loaded credential, or None if not found. + """ + if self._invocation_context.credential_service is None: + raise ValueError("Credential service is not initialized.") + return await self._invocation_context.credential_service.load_credential( + auth_config, self + ) + + def get_auth_response(self, auth_config: AuthConfig) -> AuthCredential | None: + """Gets the auth response credential from session state. + + This method retrieves an authentication credential that was previously + stored in session state after a user completed an OAuth flow or other + authentication process. + + Args: + auth_config: The authentication configuration for the credential. + + Returns: + The auth credential from the auth response, or None if not found. + """ + from ..auth.auth_handler import AuthHandler + + return AuthHandler(auth_config).get_auth_response(self.state) + + def request_credential(self, auth_config: AuthConfig) -> None: + """Requests a credential for the current tool call. + + This method can only be called in a tool context where function_call_id + is set. For callback contexts, use save_credential/load_credential instead. + + Args: + auth_config: The authentication configuration for the credential. + + Raises: + ValueError: If function_call_id is not set. + """ + from ..auth.auth_handler import AuthHandler + + if not self.function_call_id: + raise ValueError( + "request_credential requires function_call_id. " + "This method can only be used in a tool context, not a callback " + "context. Consider using save_credential/load_credential instead." + ) + self._event_actions.requested_auth_configs[self.function_call_id] = ( + AuthHandler(auth_config).generate_auth_request() + ) + + # ============================================================================ + # Tool methods + # ============================================================================ + + def request_confirmation( + self, + *, + hint: str | None = None, + payload: Any | None = None, + ) -> None: + """Requests confirmation for the current tool call. + + This method can only be called in a tool context where function_call_id + is set. + + Args: + hint: A hint to the user on how to confirm the tool call. + payload: The payload used to confirm the tool call. + + Raises: + ValueError: If function_call_id is not set. + """ + from ..tools.tool_confirmation import ToolConfirmation + + if not self.function_call_id: + raise ValueError( + "request_confirmation requires function_call_id. " + "This method can only be used in a tool context." + ) + self._event_actions.requested_tool_confirmations[self.function_call_id] = ( + ToolConfirmation( + hint=hint, + payload=payload, + ) + ) + + # ============================================================================ + # Memory methods + # ============================================================================ + + async def add_session_to_memory(self) -> None: + """Triggers memory generation for the current session. + + This method saves the current session's events to the memory service, + enabling the agent to recall information from past interactions. + + Raises: + ValueError: If memory service is not available. + + Example: + ```python + async def my_after_agent_callback(ctx: Context): + # Save conversation to memory at the end of each interaction + await ctx.add_session_to_memory() + ``` + """ + if self._invocation_context.memory_service is None: + raise ValueError( + "Cannot add session to memory: memory service is not available." + ) + await self._invocation_context.memory_service.add_session_to_memory( + self._invocation_context.session + ) + + async def add_events_to_memory( + self, + *, + events: Sequence[Event], + custom_metadata: Mapping[str, object] | None = None, + ) -> None: + """Adds an explicit list of events to the memory service. + + Uses this callback's current session identifiers as memory scope. + + Args: + events: Explicit events to add to memory. + custom_metadata: Optional standard metadata for memory generation. + + Raises: + ValueError: If memory service is not available. + """ + if self._invocation_context.memory_service is None: + raise ValueError( + "Cannot add events to memory: memory service is not available." + ) + await self._invocation_context.memory_service.add_events_to_memory( + app_name=self._invocation_context.session.app_name, + user_id=self._invocation_context.session.user_id, + session_id=self._invocation_context.session.id, + events=events, + custom_metadata=custom_metadata, + ) + + async def search_memory(self, query: str) -> SearchMemoryResponse: + """Searches the memory of the current user. + + Args: + query: The search query. + + Returns: + The search results from the memory service. + + Raises: + ValueError: If memory service is not available. + """ + if self._invocation_context.memory_service is None: + raise ValueError("Memory service is not available.") + return await self._invocation_context.memory_service.search_memory( + app_name=self._invocation_context.app_name, + user_id=self._invocation_context.user_id, + query=query, + ) diff --git a/tests/unittests/agents/test_context.py b/tests/unittests/agents/test_context.py new file mode 100644 index 00000000..3abff822 --- /dev/null +++ b/tests/unittests/agents/test_context.py @@ -0,0 +1,488 @@ +# 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.mock import AsyncMock +from unittest.mock import MagicMock +from unittest.mock import patch + +from google.adk.agents.context import Context +from google.adk.auth import auth_handler +from google.adk.auth.auth_credential import AuthCredential +from google.adk.auth.auth_credential import AuthCredentialTypes +from google.adk.auth.auth_tool import AuthConfig +from google.adk.memory.base_memory_service import SearchMemoryResponse +from google.adk.tools.tool_confirmation import ToolConfirmation +from google.genai.types import Part +import pytest + + +@pytest.fixture +def mock_invocation_context(): + """Create a mock invocation context for testing.""" + mock_context = MagicMock() + mock_context.invocation_id = "test-invocation-id" + mock_context.agent.name = "test-agent-name" + mock_context.session.state = {"key1": "value1", "key2": "value2"} + mock_context.session.id = "test-session-id" + mock_context.app_name = "test-app" + mock_context.user_id = "test-user" + mock_context.artifact_service = None + mock_context.credential_service = None + mock_context.memory_service = None + return mock_context + + +@pytest.fixture +def mock_artifact_service(): + """Create a mock artifact service for testing.""" + mock_service = AsyncMock() + mock_service.list_artifact_keys.return_value = [ + "file1.txt", + "file2.txt", + "file3.txt", + ] + return mock_service + + +@pytest.fixture +def mock_auth_config(mocker): + """Create a mock auth config for testing.""" + return mocker.create_autospec(AuthConfig, instance=True) + + +@pytest.fixture +def mock_auth_credential(mocker): + """Create a mock auth credential for testing.""" + mock_credential = mocker.create_autospec(AuthCredential, instance=True) + mock_credential.auth_type = AuthCredentialTypes.OAUTH2 + return mock_credential + + +class TestContextInitialization: + """Test Context initialization.""" + + def test_initialization_without_function_call_id( + self, mock_invocation_context + ): + """Test Context initialization without function_call_id.""" + context = Context(mock_invocation_context) + + assert context._invocation_context == mock_invocation_context + assert context._event_actions is not None + assert context._state is not None + assert context.function_call_id is None + assert context.tool_confirmation is None + + def test_initialization_with_function_call_id(self, mock_invocation_context): + """Test Context initialization with function_call_id.""" + context = Context( + mock_invocation_context, + function_call_id="test-function-call-id", + ) + + assert context.function_call_id == "test-function-call-id" + assert context.tool_confirmation is None + + def test_initialization_with_tool_confirmation(self, mock_invocation_context): + """Test Context initialization with tool_confirmation.""" + tool_confirmation = ToolConfirmation( + hint="test hint", payload={"key": "value"} + ) + context = Context( + mock_invocation_context, + function_call_id="test-function-call-id", + tool_confirmation=tool_confirmation, + ) + + assert context.function_call_id == "test-function-call-id" + assert context.tool_confirmation == tool_confirmation + assert context.tool_confirmation.hint == "test hint" + assert context.tool_confirmation.payload == {"key": "value"} + + def test_state_property(self, mock_invocation_context): + """Test that state property returns mutable state.""" + context = Context(mock_invocation_context) + + assert context.state["key1"] == "value1" + assert context.state["key2"] == "value2" + + def test_actions_property(self, mock_invocation_context): + """Test that actions property returns event_actions.""" + context = Context(mock_invocation_context) + + assert context.actions is context._event_actions + + +class TestContextListArtifacts: + """Test the list_artifacts method in Context.""" + + async def test_list_artifacts_returns_artifact_keys( + self, mock_invocation_context, mock_artifact_service + ): + """Test that list_artifacts returns the artifact keys from the service.""" + mock_invocation_context.artifact_service = mock_artifact_service + context = Context(mock_invocation_context) + + result = await context.list_artifacts() + + assert result == ["file1.txt", "file2.txt", "file3.txt"] + mock_artifact_service.list_artifact_keys.assert_called_once_with( + app_name="test-app", + user_id="test-user", + session_id="test-session-id", + ) + + async def test_list_artifacts_raises_value_error_when_service_is_none( + self, mock_invocation_context + ): + """Test that list_artifacts raises ValueError when no artifact service.""" + mock_invocation_context.artifact_service = None + context = Context(mock_invocation_context) + + with pytest.raises( + ValueError, match="Artifact service is not initialized." + ): + await context.list_artifacts() + + +class TestContextSaveLoadArtifact: + """Test save_artifact and load_artifact methods in Context.""" + + async def test_save_artifact(self, mock_invocation_context): + """Test save_artifact method.""" + artifact_service = AsyncMock() + artifact_service.save_artifact.return_value = 1 + mock_invocation_context.artifact_service = artifact_service + + context = Context(mock_invocation_context) + test_artifact = Part.from_text(text="test content") + + version = await context.save_artifact("test_file.txt", test_artifact) + + artifact_service.save_artifact.assert_called_once_with( + app_name="test-app", + user_id="test-user", + session_id="test-session-id", + filename="test_file.txt", + artifact=test_artifact, + custom_metadata=None, + ) + assert version == 1 + assert context.actions.artifact_delta["test_file.txt"] == 1 + + async def test_load_artifact(self, mock_invocation_context): + """Test load_artifact method.""" + artifact_service = AsyncMock() + test_artifact = Part.from_text(text="test content") + artifact_service.load_artifact.return_value = test_artifact + mock_invocation_context.artifact_service = artifact_service + + context = Context(mock_invocation_context) + + result = await context.load_artifact("test_file.txt") + + artifact_service.load_artifact.assert_called_once_with( + app_name="test-app", + user_id="test-user", + session_id="test-session-id", + filename="test_file.txt", + version=None, + ) + assert result == test_artifact + + async def test_load_artifact_with_version(self, mock_invocation_context): + """Test load_artifact method with specific version.""" + artifact_service = AsyncMock() + test_artifact = Part.from_text(text="test content") + artifact_service.load_artifact.return_value = test_artifact + mock_invocation_context.artifact_service = artifact_service + + context = Context(mock_invocation_context) + + result = await context.load_artifact("test_file.txt", version=2) + + artifact_service.load_artifact.assert_called_once_with( + app_name="test-app", + user_id="test-user", + session_id="test-session-id", + filename="test_file.txt", + version=2, + ) + assert result == test_artifact + + +class TestContextCredentialMethods: + """Test credential methods in Context.""" + + async def test_save_credential_with_service( + self, mock_invocation_context, mock_auth_config + ): + """Test save_credential when credential service is available.""" + credential_service = AsyncMock() + mock_invocation_context.credential_service = credential_service + + context = Context(mock_invocation_context) + await context.save_credential(mock_auth_config) + + credential_service.save_credential.assert_called_once_with( + mock_auth_config, context + ) + + async def test_save_credential_no_service( + self, mock_invocation_context, mock_auth_config + ): + """Test save_credential when credential service is not available.""" + mock_invocation_context.credential_service = None + + context = Context(mock_invocation_context) + + with pytest.raises( + ValueError, match="Credential service is not initialized" + ): + await context.save_credential(mock_auth_config) + + async def test_load_credential_with_service( + self, mock_invocation_context, mock_auth_config, mock_auth_credential + ): + """Test load_credential when credential service is available.""" + credential_service = AsyncMock() + credential_service.load_credential.return_value = mock_auth_credential + mock_invocation_context.credential_service = credential_service + + context = Context(mock_invocation_context) + result = await context.load_credential(mock_auth_config) + + credential_service.load_credential.assert_called_once_with( + mock_auth_config, context + ) + assert result == mock_auth_credential + + async def test_load_credential_no_service( + self, mock_invocation_context, mock_auth_config + ): + """Test load_credential when credential service is not available.""" + mock_invocation_context.credential_service = None + + context = Context(mock_invocation_context) + + with pytest.raises( + ValueError, match="Credential service is not initialized" + ): + await context.load_credential(mock_auth_config) + + +class TestContextGetAuthResponse: + """Test get_auth_response method in Context.""" + + def test_get_auth_response(self, mock_invocation_context, mock_auth_config): + """Test get_auth_response method.""" + context = Context(mock_invocation_context) + + with patch.object( + auth_handler, "AuthHandler", autospec=True + ) as mock_auth_handler: + mock_handler_instance = mock_auth_handler.return_value + mock_handler_instance.get_auth_response.return_value = "auth-response" + + result = context.get_auth_response(mock_auth_config) + + mock_auth_handler.assert_called_once_with(mock_auth_config) + mock_handler_instance.get_auth_response.assert_called_once_with( + context.state + ) + assert result == "auth-response" + + +class TestContextRequestCredential: + """Test request_credential method in Context.""" + + def test_request_credential_with_function_call_id( + self, mock_invocation_context, mock_auth_config + ): + """Test request_credential when function_call_id is set.""" + context = Context( + mock_invocation_context, + function_call_id="test-function-call-id", + ) + + with patch.object( + auth_handler, "AuthHandler", autospec=True + ) as mock_auth_handler: + mock_handler_instance = mock_auth_handler.return_value + mock_handler_instance.generate_auth_request.return_value = "auth-request" + + context.request_credential(mock_auth_config) + + mock_auth_handler.assert_called_once_with(mock_auth_config) + mock_handler_instance.generate_auth_request.assert_called_once() + assert ( + context.actions.requested_auth_configs["test-function-call-id"] + == "auth-request" + ) + + def test_request_credential_without_function_call_id_raises( + self, mock_invocation_context, mock_auth_config + ): + """Test request_credential raises ValueError when no function_call_id.""" + context = Context(mock_invocation_context) + + with pytest.raises( + ValueError, + match="request_credential requires function_call_id", + ): + context.request_credential(mock_auth_config) + + +class TestContextRequestConfirmation: + """Test request_confirmation method in Context.""" + + def test_request_confirmation_with_function_call_id( + self, mock_invocation_context + ): + """Test request_confirmation when function_call_id is set.""" + context = Context( + mock_invocation_context, + function_call_id="test-function-call-id", + ) + + context.request_confirmation( + hint="Please confirm", payload={"action": "delete"} + ) + + confirmation = context.actions.requested_tool_confirmations[ + "test-function-call-id" + ] + assert confirmation.hint == "Please confirm" + assert confirmation.payload == {"action": "delete"} + + def test_request_confirmation_with_only_hint(self, mock_invocation_context): + """Test request_confirmation with only hint provided.""" + context = Context( + mock_invocation_context, + function_call_id="test-function-call-id", + ) + + context.request_confirmation(hint="Confirm this action") + + confirmation = context.actions.requested_tool_confirmations[ + "test-function-call-id" + ] + assert confirmation.hint == "Confirm this action" + assert confirmation.payload is None + + def test_request_confirmation_without_function_call_id_raises( + self, mock_invocation_context + ): + """Test request_confirmation raises ValueError when no function_call_id.""" + context = Context(mock_invocation_context) + + with pytest.raises( + ValueError, + match="request_confirmation requires function_call_id", + ): + context.request_confirmation() + + +class TestContextMemoryMethods: + """Test memory methods in Context.""" + + async def test_add_session_to_memory_success(self, mock_invocation_context): + """Test that add_session_to_memory calls the memory service correctly.""" + memory_service = AsyncMock() + mock_invocation_context.memory_service = memory_service + + context = Context(mock_invocation_context) + await context.add_session_to_memory() + + memory_service.add_session_to_memory.assert_called_once_with( + mock_invocation_context.session + ) + + async def test_add_session_to_memory_no_service_raises( + self, mock_invocation_context + ): + """Test that add_session_to_memory raises ValueError when memory service is None.""" + mock_invocation_context.memory_service = None + + context = Context(mock_invocation_context) + + with pytest.raises( + ValueError, + match=( + r"Cannot add session to memory: memory service is not available\." + ), + ): + await context.add_session_to_memory() + + async def test_search_memory_success(self, mock_invocation_context, mocker): + """Test that search_memory calls the memory service correctly.""" + memory_service = AsyncMock() + mock_search_response = mocker.create_autospec( + SearchMemoryResponse, instance=True + ) + memory_service.search_memory.return_value = mock_search_response + mock_invocation_context.memory_service = memory_service + + context = Context(mock_invocation_context) + result = await context.search_memory("test query") + + memory_service.search_memory.assert_called_once_with( + app_name="test-app", + user_id="test-user", + query="test query", + ) + assert result == mock_search_response + + async def test_search_memory_no_service_raises(self, mock_invocation_context): + """Test that search_memory raises ValueError when memory service is None.""" + mock_invocation_context.memory_service = None + + context = Context(mock_invocation_context) + + with pytest.raises(ValueError, match="Memory service is not available."): + await context.search_memory("test query") + + async def test_add_events_to_memory_success(self, mock_invocation_context): + """Test that add_events_to_memory calls the memory service correctly.""" + memory_service = AsyncMock() + mock_invocation_context.memory_service = memory_service + test_event = MagicMock() + + context = Context(mock_invocation_context) + await context.add_events_to_memory( + events=[test_event], + custom_metadata={"ttl": "6000s"}, + ) + + memory_service.add_events_to_memory.assert_called_once_with( + app_name=mock_invocation_context.session.app_name, + user_id=mock_invocation_context.session.user_id, + session_id=mock_invocation_context.session.id, + events=[test_event], + custom_metadata={"ttl": "6000s"}, + ) + + async def test_add_events_to_memory_no_service_raises( + self, mock_invocation_context + ): + """Test that add_events_to_memory raises ValueError when no service.""" + mock_invocation_context.memory_service = None + + context = Context(mock_invocation_context) + + with pytest.raises( + ValueError, + match=r"Cannot add events to memory: memory service is not available\.", + ): + await context.add_events_to_memory(events=[MagicMock()])