mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
chore: Unify CallbackContext and ToolContext into the Context class
Co-authored-by: Xuan Yang <xygoogle@google.com> PiperOrigin-RevId: 869400758
This commit is contained in:
Xuan Yang
committed by
Copybara-Service
Copybara-Service
parent
4f71b45425
commit
647d3b16a3
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user