feat: Define an AgentState to be used for resuming agent invocation

PiperOrigin-RevId: 811414736
This commit is contained in:
Shangjie Chen
2025-09-25 10:49:49 -07:00
committed by Copybara-Service
parent 1589fcdd86
commit 839d2e43bb
4 changed files with 41 additions and 2 deletions
+9
View File
@@ -66,6 +66,15 @@ AfterAgentCallback: TypeAlias = Union[
SelfAgent = TypeVar('SelfAgent', bound='BaseAgent')
@experimental
class BaseAgentState(BaseModel):
"""Base class for all agent states."""
model_config = ConfigDict(
extra='forbid',
)
class BaseAgent(BaseModel):
"""Base class for all agents in Agent Development Kit."""
+12 -1
View File
@@ -21,7 +21,6 @@ from typing import AsyncGenerator
from typing import ClassVar
from typing import Dict
from typing import Optional
from typing import Type
from typing_extensions import override
@@ -30,10 +29,22 @@ from ..events.event import Event
from ..utils.context_utils import Aclosing
from ..utils.feature_decorator import experimental
from .base_agent import BaseAgent
from .base_agent import BaseAgentState
from .base_agent_config import BaseAgentConfig
from .loop_agent_config import LoopAgentConfig
@experimental
class LoopAgentState(BaseAgentState):
"""State for LoopAgent."""
current_sub_agent: str = ''
"""The name of the current sub-agent to run in the loop."""
times_looped: int = 0
"""The number of times the loop agent has looped."""
class LoopAgent(BaseAgent):
"""A shell agent that run its sub-agents in a loop.
+11 -1
View File
@@ -24,13 +24,23 @@ from typing_extensions import override
from ..events.event import Event
from ..utils.context_utils import Aclosing
from ..utils.feature_decorator import experimental
from .base_agent import BaseAgent
from .base_agent import BaseAgentConfig
from .base_agent import BaseAgentState
from .base_agent_config import BaseAgentConfig
from .invocation_context import InvocationContext
from .llm_agent import LlmAgent
from .sequential_agent_config import SequentialAgentConfig
@experimental
class SequentialAgentState(BaseAgentState):
"""State for SequentialAgent."""
current_sub_agent: str = ''
"""The name of the current sub-agent to run."""
class SequentialAgent(BaseAgent):
"""A shell agent that runs its sub-agents in sequence."""
+9
View File
@@ -14,6 +14,7 @@
from __future__ import annotations
from typing import Any
from typing import Optional
from google.genai.types import Content
@@ -95,3 +96,11 @@ class EventActions(BaseModel):
compaction: Optional[EventCompaction] = None
"""The compaction of the events."""
end_of_agent: Optional[bool] = None
"""If true, the current agent has finished its current run. Note that there
can be multiple events with end_of_agent=True for the same agent within one
invocation when there is a loop."""
agent_state: Optional[dict[str, Any]] = None
"""The agent state at the current event."""