feat: Add core checkpointing primitive for base agent

PiperOrigin-RevId: 811458903

Author: DeanChensj

src/google/adk/agents/base_agent.py

@@ -38,6 +38,7 @@
 from typing_extensions import TypeAlias
 
 from ..events.event import Event
+from ..events.event_actions import EventActions
 from ..telemetry import tracing
 from ..telemetry.tracing import tracer
 from ..utils.context_utils import Aclosing
@@ -75,6 +76,9 @@ class BaseAgentState(BaseModel):
   )
 
 
+AgentState = TypeVar('AgentState', bound=BaseAgentState)
+
+
 class BaseAgent(BaseModel):
   """Base class for all agents in Agent Development Kit."""
 
@@ -155,6 +159,57 @@ class MyAgent(BaseAgent):
       response and appended to event history as agent response.
   """
 
+  def _load_agent_state(
+      self,
+      ctx: InvocationContext,
+      state_type: Type[AgentState],
+      default_state: AgentState,
+  ) -> tuple[AgentState, bool]:
+    """Loads the agent state from the invocation context, handling resumption.
+
+    Args:
+      ctx: The invocation context.
+      state_type: The type of the agent state.
+      default_state: The default state to use if not resuming.
+
+    Returns:
+        tuple[AgentState, bool]: The current state and a boolean indicating if
+        resuming.
+    """
+    if self.name not in ctx.agent_states:
+      return default_state, False
+    else:
+      return state_type.model_validate(ctx.agent_states.get(self.name)), True
+
+  def _create_agent_state_event(
+      self,
+      ctx: InvocationContext,
+      *,
+      state: Optional[BaseAgentState] = None,
+      end_of_agent: bool = False,
+  ) -> Event:
+    """Creates an event for agent state.
+
+    Args:
+      ctx: The invocation context.
+      state: The agent state to checkpoint.
+      end_of_agent: Whether the agent is finished running.
+
+    Returns:
+      An Event object representing the checkpoint.
+    """
+    event_actions = EventActions()
+    if state:
+      event_actions.agent_state = state.model_dump(mode='json')
+    if end_of_agent:
+      event_actions.end_of_agent = True
+    return Event(
+        invocation_id=ctx.invocation_id,
+        author=self.name,
+        branch=ctx.branch,
+        actions=event_actions,
+    )
+
   def clone(
       self: SelfAgent, update: Mapping[str, Any] | None = None
   ) -> SelfAgent:

src/google/adk/agents/invocation_context.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from typing import Any
 from typing import Optional
 import uuid
 
@@ -162,6 +163,12 @@ class InvocationContext(BaseModel):
   session: Session
   """The current session of this invocation context. Readonly."""
 
+  agent_states: dict[str, dict[str, Any]] = Field(default_factory=dict)
+  """The state of the agent for this invocation."""
+
+  end_of_agents: dict[str, bool] = Field(default_factory=dict)
+  """The end of agent status for each agent in this invocation."""
+
   end_invocation: bool = False
   """Whether to end this invocation.
 
@@ -201,6 +208,11 @@ class InvocationContext(BaseModel):
   of this invocation.
   """
 
+  def reset_agent_state(self, agent_name: str) -> None:
+    """Resets the state of an agent, allowing it to be re-run."""
+    self.agent_states.pop(agent_name, None)
+    self.end_of_agents.pop(agent_name, None)
+
   def increment_llm_call_count(
       self,
   ):

tests/unittests/agents/test_base_agent.py

@@ -23,6 +23,7 @@
 from unittest import mock
 
 from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.base_agent import BaseAgentState
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.events.event import Event
@@ -732,6 +733,39 @@ async def test_run_live_incomplete_agent(request: pytest.FixtureRequest):
     [e async for e in agent.run_live(parent_ctx)]
 
 
+@pytest.mark.asyncio
+async def test_create_agent_state_event(request: pytest.FixtureRequest):
+  # Arrange
+  agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
+  ctx = await _create_parent_invocation_context(
+      request.function.__name__, agent, branch='test_branch'
+  )
+  state = BaseAgentState()
+
+  # Act
+  event = agent._create_agent_state_event(ctx, state=state)
+
+  # Assert
+  assert event.invocation_id == ctx.invocation_id
+  assert event.author == agent.name
+  assert event.branch == 'test_branch'
+  assert event.actions is not None
+  assert event.actions.agent_state is not None
+  assert event.actions.agent_state == state.model_dump(mode='json')
+  assert not event.actions.end_of_agent
+
+  # Act
+  event = agent._create_agent_state_event(ctx, end_of_agent=True)
+
+  # Assert
+  assert event.invocation_id == ctx.invocation_id
+  assert event.author == agent.name
+  assert event.branch == 'test_branch'
+  assert event.actions is not None
+  assert event.actions.end_of_agent
+  assert event.actions.agent_state is None
+
+
 def test_set_parent_agent_for_sub_agents(request: pytest.FixtureRequest):
   sub_agents: list[BaseAgent] = [
       _TestingAgent(name=f'{request.function.__name__}_sub_agent_1'),
@@ -854,3 +888,54 @@ def test_set_parent_agent_for_sub_agent_twice(
 
 if __name__ == '__main__':
   pytest.main([__file__])
+
+
+class _TestAgentState(BaseAgentState):
+  test_field: str = ''
+
+
+@pytest.mark.asyncio
+async def test_load_agent_state_no_resume():
+  agent = BaseAgent(name='test_agent')
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  ctx = InvocationContext(
+      invocation_id='test_invocation',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+  )
+  default_state = _TestAgentState(test_field='default')
+
+  state, is_resuming = agent._load_agent_state(
+      ctx, _TestAgentState, default_state
+  )
+
+  assert not is_resuming
+  assert state == default_state
+
+
+@pytest.mark.asyncio
+async def test_load_agent_state_with_resume():
+  agent = BaseAgent(name='test_agent')
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  ctx = InvocationContext(
+      invocation_id='test_invocation',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+  )
+  persisted_state = _TestAgentState(test_field='resumed')
+  ctx.agent_states[agent.name] = persisted_state.model_dump(mode='json')
+
+  state, is_resuming = agent._load_agent_state(
+      ctx, _TestAgentState, _TestAgentState()
+  )
+
+  assert is_resuming
+  assert state == persisted_state