From 647d3b16a3ae2ddb5410b78b25b7494bf4c51681 Mon Sep 17 00:00:00 2001 From: Xuan Yang Date: Thu, 12 Feb 2026 15:01:15 -0800 Subject: [PATCH] chore: Unify CallbackContext and ToolContext into the Context class Co-authored-by: Xuan Yang PiperOrigin-RevId: 869400758 --- src/google/adk/agents/callback_context.py | 239 +----------------- src/google/adk/agents/context.py | 10 + src/google/adk/tools/tool_context.py | 92 +------ .../test_tool_auth_handler.py | 2 +- 4 files changed, 23 insertions(+), 320 deletions(-) diff --git a/src/google/adk/agents/callback_context.py b/src/google/adk/agents/callback_context.py index d733540b..18b15920 100644 --- a/src/google/adk/agents/callback_context.py +++ b/src/google/adk/agents/callback_context.py @@ -14,240 +14,9 @@ from __future__ import annotations -from collections.abc import Mapping -from collections.abc import Sequence -from typing import Any -from typing import Optional -from typing import TYPE_CHECKING - -from typing_extensions import override - +from .context import Context +# Keep ReadonlyContext for backward compatibility 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 ..sessions.state import State - from .invocation_context import InvocationContext - - -class CallbackContext(ReadonlyContext): - """The context of various callbacks within an agent run.""" - - def __init__( - self, - invocation_context: InvocationContext, - *, - event_actions: Optional[EventActions] = None, - ) -> None: - 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, - ) - - @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 - - async def load_artifact( - self, filename: str, version: Optional[int] = None - ) -> Optional[types.Part]: - """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: Optional[dict[str, Any]] = 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: Optional[int] = None - ) -> Optional[ArtifactVersion]: - """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, - ) - - 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 - ) -> Optional[AuthCredential]: - """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 - ) -> Optional[AuthCredential]: - """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) - - 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(callback_context: CallbackContext): - # Save conversation to memory at the end of each interaction - await callback_context.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, - ) +# CallbackContext is unified into Context +CallbackContext = Context diff --git a/src/google/adk/agents/context.py b/src/google/adk/agents/context.py index 79942681..1ec9a5f8 100644 --- a/src/google/adk/agents/context.py +++ b/src/google/adk/agents/context.py @@ -76,11 +76,21 @@ class Context(ReadonlyContext): """The function call id of the current tool call.""" return self._function_call_id + @function_call_id.setter + def function_call_id(self, value: str | None) -> None: + """Sets the function call id of the current tool call.""" + self._function_call_id = value + @property def tool_confirmation(self) -> ToolConfirmation | None: """The tool confirmation of the current tool call.""" return self._tool_confirmation + @tool_confirmation.setter + def tool_confirmation(self, value: ToolConfirmation | None) -> None: + """Sets the tool confirmation of the current tool call.""" + self._tool_confirmation = value + @property @override def state(self) -> State: diff --git a/src/google/adk/tools/tool_context.py b/src/google/adk/tools/tool_context.py index 12ae55db..6d645c66 100644 --- a/src/google/adk/tools/tool_context.py +++ b/src/google/adk/tools/tool_context.py @@ -14,93 +14,17 @@ from __future__ import annotations -from typing import Any -from typing import Optional -from typing import TYPE_CHECKING - +# Keep CallbackContext for backward compatibility from ..agents.callback_context import CallbackContext +from ..agents.context import Context +# Keep AuthCredential for backward compatibility from ..auth.auth_credential import AuthCredential +# Keep AuthHandler for backward compatibility from ..auth.auth_handler import AuthHandler +# Keep AuthConfig for backward compatibility from ..auth.auth_tool import AuthConfig +# Keep ToolConfirmation for backward compatibility from .tool_confirmation import ToolConfirmation -if TYPE_CHECKING: - from ..agents.invocation_context import InvocationContext - from ..events.event_actions import EventActions - from ..memory.base_memory_service import SearchMemoryResponse - - -class ToolContext(CallbackContext): - """The context of the tool. - - This class provides the context for a tool invocation, including access to - the invocation context, function call ID, event actions, and authentication - response. It also provides methods for requesting credentials, retrieving - authentication responses, listing artifacts, and searching memory. - - Attributes: - invocation_context: The invocation context of the tool. - function_call_id: The function call id of the current tool call. This id was - returned in the function call event from LLM to identify a function call. - If LLM didn't return this id, ADK will assign one to it. This id is used - to map function call response to the original function call. - event_actions: The event actions of the current tool call. - tool_confirmation: The tool confirmation of the current tool call. - """ - - def __init__( - self, - invocation_context: InvocationContext, - *, - function_call_id: Optional[str] = None, - event_actions: Optional[EventActions] = None, - tool_confirmation: Optional[ToolConfirmation] = None, - ): - super().__init__(invocation_context, event_actions=event_actions) - self.function_call_id = function_call_id - self.tool_confirmation = tool_confirmation - - @property - def actions(self) -> EventActions: - return self._event_actions - - def request_credential(self, auth_config: AuthConfig) -> None: - if not self.function_call_id: - raise ValueError('function_call_id is not set.') - self._event_actions.requested_auth_configs[self.function_call_id] = ( - AuthHandler(auth_config).generate_auth_request() - ) - - def get_auth_response(self, auth_config: AuthConfig) -> AuthCredential: - return AuthHandler(auth_config).get_auth_response(self.state) - - def request_confirmation( - self, - *, - hint: Optional[str] = None, - payload: Optional[Any] = None, - ) -> None: - """Requests confirmation for the given function call. - - Args: - hint: A hint to the user on how to confirm the tool call. - payload: The payload used to confirm the tool call. - """ - if not self.function_call_id: - raise ValueError('function_call_id is not set.') - self._event_actions.requested_tool_confirmations[self.function_call_id] = ( - ToolConfirmation( - hint=hint, - payload=payload, - ) - ) - - async def search_memory(self, query: str) -> SearchMemoryResponse: - """Searches the memory of the current user.""" - 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, - ) +# ToolContext is unified into Context +ToolContext = Context diff --git a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py index 35914c46..a6babce6 100644 --- a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py +++ b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_tool_auth_handler.py @@ -172,7 +172,7 @@ async def test_openid_connect_with_auth_response( oauth2=OAuth2Auth(auth_response_uri='test_auth_response_uri'), ) mock_auth_handler.get_auth_response.return_value = returned_credential - mock_auth_handler_path = 'google.adk.tools.tool_context.AuthHandler' + mock_auth_handler_path = 'google.adk.auth.auth_handler.AuthHandler' monkeypatch.setattr( mock_auth_handler_path, lambda *args, **kwargs: mock_auth_handler )