change AgentBase from abstract class to protocol + remove invocation_state and structured_output_model

Author: awsarron

src/strands/agent/agent.py

@@ -64,7 +64,6 @@
 from ..types.tools import ToolResult, ToolUse
 from ..types.traces import AttributeValue
 from .agent_result import AgentResult
-from .base import AgentBase
 from .conversation_manager import (
     ConversationManager,
     SlidingWindowConversationManager,
@@ -89,7 +88,7 @@ class _DefaultCallbackHandlerSentinel:
 _DEFAULT_AGENT_ID = "default"
 
 
-class Agent(AgentBase):
+class Agent:
     """Core Agent implementation.
 
     An agent orchestrates the following workflow:

src/strands/agent/base.py

@@ -3,115 +3,74 @@
 Defines the minimal interface that all agent types must implement.
 """
 
-from abc import ABC, abstractmethod
-from typing import Any, AsyncIterator, Type
-
-from pydantic import BaseModel
+from typing import Any, AsyncIterator, Protocol, runtime_checkable
 
 from ..types.agent import AgentInput
 from .agent_result import AgentResult
 from .state import AgentState
 
 
-class AgentBase(ABC):
-    """Abstract interface for all agent types in Strands.
+@runtime_checkable
+class AgentBase(Protocol):
+    """Protocol defining the interface for all agent types in Strands.
 
-    This interface defines the minimal contract that all agent implementations
+    This protocol defines the minimal contract that all agent implementations
     must satisfy.
     """
 
-    @property
-    @abstractmethod
-    def agent_id(self) -> str:
-        """Unique identifier for the agent.
-
-        Returns:
-            Unique string identifier for this agent instance.
-        """
-        pass
-
-    @property
-    @abstractmethod
-    def name(self) -> str:
-        """Human-readable name of the agent.
+    agent_id: str
+    """Unique identifier for the agent."""
 
-        Returns:
-            Display name for the agent.
-        """
-        pass
-
-    @property
-    @abstractmethod
-    def state(self) -> AgentState:
-        """Current state of the agent.
+    name: str
+    """Human-readable name of the agent."""
 
-        Returns:
-            AgentState object containing stateful information.
-        """
-        pass
+    state: AgentState
+    """Current state of the agent."""
 
-    @abstractmethod
     async def invoke_async(
         self,
         prompt: AgentInput = None,
-        *,
-        invocation_state: dict[str, Any] | None = None,
-        structured_output_model: Type[BaseModel] | None = None,
         **kwargs: Any,
     ) -> AgentResult:
         """Asynchronously invoke the agent with the given prompt.
 
         Args:
             prompt: Input to the agent.
-            invocation_state: Optional state to pass to the agent invocation.
-            structured_output_model: Optional Pydantic model for structured output.
-            **kwargs: Additional provider-specific arguments.
+            **kwargs: Additional arguments.
 
         Returns:
             AgentResult containing the agent's response.
         """
-        pass
+        ...
 
-    @abstractmethod
     def __call__(
         self,
         prompt: AgentInput = None,
-        *,
-        invocation_state: dict[str, Any] | None = None,
-        structured_output_model: Type[BaseModel] | None = None,
         **kwargs: Any,
     ) -> AgentResult:
         """Synchronously invoke the agent with the given prompt.
 
         Args:
             prompt: Input to the agent.
-            invocation_state: Optional state to pass to the agent invocation.
-            structured_output_model: Optional Pydantic model for structured output.
-            **kwargs: Additional provider-specific arguments.
+            **kwargs: Additional arguments.
 
         Returns:
             AgentResult containing the agent's response.
         """
-        pass
+        ...
 
-    @abstractmethod
     def stream_async(
         self,
         prompt: AgentInput = None,
-        *,
-        invocation_state: dict[str, Any] | None = None,
-        structured_output_model: Type[BaseModel] | None = None,
         **kwargs: Any,
     ) -> AsyncIterator[Any]:
         """Stream agent execution asynchronously.
 
         Args:
             prompt: Input to the agent.
-            invocation_state: Optional state to pass to the agent invocation.
-            structured_output_model: Optional Pydantic model for structured output.
-            **kwargs: Additional provider-specific arguments.
+            **kwargs: Additional arguments.
 
         Yields:
             Events representing the streaming execution.
         """
-        pass
+        ...

src/strands/multiagent/graph.py

@@ -24,7 +24,7 @@
 from opentelemetry import trace as trace_api
 
 from .._async import run_async
-from ..agent import Agent
+from ..agent import Agent, AgentBase
 from ..agent.state import AgentState
 from ..experimental.hooks.multiagent import (
     AfterMultiAgentInvocationEvent,
@@ -154,7 +154,7 @@ class GraphNode:
     """
 
     node_id: str
-    executor: Agent | MultiAgentBase
+    executor: AgentBase | MultiAgentBase
     dependencies: set["GraphNode"] = field(default_factory=set)
     execution_status: Status = Status.PENDING
     result: NodeResult | None = None
@@ -199,7 +199,7 @@ def __eq__(self, other: Any) -> bool:
 
 
 def _validate_node_executor(
-    executor: Agent | MultiAgentBase, existing_nodes: dict[str, GraphNode] | None = None
+    executor: AgentBase | MultiAgentBase, existing_nodes: dict[str, GraphNode] | None = None
 ) -> None:
     """Validate a node executor for graph compatibility.
 
@@ -238,8 +238,8 @@ def __init__(self) -> None:
         self._session_manager: Optional[SessionManager] = None
         self._hooks: Optional[list[HookProvider]] = None
 
-    def add_node(self, executor: Agent | MultiAgentBase, node_id: str | None = None) -> GraphNode:
-        """Add an Agent or MultiAgentBase instance as a node to the graph."""
+    def add_node(self, executor: AgentBase | MultiAgentBase, node_id: str | None = None) -> GraphNode:
+        """Add an AgentBase or MultiAgentBase instance as a node to the graph."""
         _validate_node_executor(executor, self.nodes)
 
         # Auto-generate node_id if not provided

src/strands/multiagent/swarm.py

@@ -23,7 +23,7 @@
 from opentelemetry import trace as trace_api
 
 from .._async import run_async
-from ..agent import Agent
+from ..agent import Agent, AgentBase
 from ..agent.state import AgentState
 from ..experimental.hooks.multiagent import (
     AfterMultiAgentInvocationEvent,
@@ -57,7 +57,7 @@ class SwarmNode:
     """Represents a node (e.g. Agent) in the swarm."""
 
     node_id: str
-    executor: Agent
+    executor: AgentBase
     _initial_messages: Messages = field(default_factory=list, init=False)
     _initial_state: AgentState = field(default_factory=AgentState, init=False)
 
@@ -212,9 +212,9 @@ class Swarm(MultiAgentBase):
 
     def __init__(
         self,
-        nodes: list[Agent],
+        nodes: list[AgentBase],
         *,
-        entry_point: Agent | None = None,
+        entry_point: AgentBase | None = None,
         max_handoffs: int = 20,
         max_iterations: int = 20,
         execution_timeout: float = 900.0,
@@ -229,8 +229,8 @@ def __init__(
 
         Args:
             id : Unique swarm id (default: None)
-            nodes: List of nodes (e.g. Agent) to include in the swarm
-            entry_point: Agent to start with. If None, uses the first agent (default: None)
+            nodes: List of nodes (e.g. AgentBase) to include in the swarm
+            entry_point: AgentBase to start with. If None, uses the first agent (default: None)
             max_handoffs: Maximum handoffs to agents and users (default: 20)
             max_iterations: Maximum node executions within the swarm (default: 20)
             execution_timeout: Total execution timeout in seconds (default: 900.0)
@@ -425,7 +425,7 @@ async def _stream_with_timeout(
                 except asyncio.TimeoutError as err:
                     raise Exception(timeout_message) from err
 
-    def _setup_swarm(self, nodes: list[Agent]) -> None:
+    def _setup_swarm(self, nodes: list[AgentBase]) -> None:
         """Initialize swarm configuration."""
         # Validate nodes before setup
         self._validate_swarm(nodes)
@@ -467,7 +467,7 @@ def _setup_swarm(self, nodes: list[Agent]) -> None:
             first_node = next(iter(self.nodes.keys()))
             logger.debug("entry_point=<%s> | using first node as entry point", first_node)
 
-    def _validate_swarm(self, nodes: list[Agent]) -> None:
+    def _validate_swarm(self, nodes: list[AgentBase]) -> None:
         """Validate swarm structure and nodes."""
         # Check for duplicate object instances
         seen_instances = set()

src/strands/session/repository_session_manager.py

@@ -17,7 +17,7 @@
 from .session_repository import SessionRepository
 
 if TYPE_CHECKING:
-    from ..agent.agent import Agent
+    from ..agent.base import AgentBase
     from ..multiagent.base import MultiAgentBase
 
 logger = logging.getLogger(__name__)
@@ -58,12 +58,12 @@ def __init__(
         # Keep track of the latest message of each agent in case we need to redact it.
         self._latest_agent_message: dict[str, Optional[SessionMessage]] = {}
 
-    def append_message(self, message: Message, agent: "Agent", **kwargs: Any) -> None:
+    def append_message(self, message: Message, agent: "AgentBase", **kwargs: Any) -> None:
         """Append a message to the agent's session.
 
         Args:
             message: Message to add to the agent in the session
-            agent: Agent to append the message to
+            agent: AgentBase to append the message to
             **kwargs: Additional keyword arguments for future extensibility.
         """
         # Calculate the next index (0 if this is the first message, otherwise increment the previous index)
@@ -77,12 +77,12 @@ def append_message(self, message: Message, agent: "Agent", **kwargs: Any) -> Non
         self._latest_agent_message[agent.agent_id] = session_message
         self.session_repository.create_message(self.session_id, agent.agent_id, session_message)
 
-    def redact_latest_message(self, redact_message: Message, agent: "Agent", **kwargs: Any) -> None:
+    def redact_latest_message(self, redact_message: Message, agent: "AgentBase", **kwargs: Any) -> None:
         """Redact the latest message appended to the session.
 
         Args:
             redact_message: New message to use that contains the redact content
-            agent: Agent to apply the message redaction to
+            agent: AgentBase to apply the message redaction to
             **kwargs: Additional keyword arguments for future extensibility.
         """
         latest_agent_message = self._latest_agent_message[agent.agent_id]
@@ -91,23 +91,23 @@ def redact_latest_message(self, redact_message: Message, agent: "Agent", **kwarg
         latest_agent_message.redact_message = redact_message
         return self.session_repository.update_message(self.session_id, agent.agent_id, latest_agent_message)
 
-    def sync_agent(self, agent: "Agent", **kwargs: Any) -> None:
+    def sync_agent(self, agent: "AgentBase", **kwargs: Any) -> None:
         """Serialize and update the agent into the session repository.
 
         Args:
-            agent: Agent to sync to the session.
+            agent: AgentBase to sync to the session.
             **kwargs: Additional keyword arguments for future extensibility.
         """
         self.session_repository.update_agent(
             self.session_id,
             SessionAgent.from_agent(agent),
         )
 
-    def initialize(self, agent: "Agent", **kwargs: Any) -> None:
+    def initialize(self, agent: "AgentBase", **kwargs: Any) -> None:
         """Initialize an agent with a session.
 
         Args:
-            agent: Agent to initialize from the session
+            agent: AgentBase to initialize from the session
             **kwargs: Additional keyword arguments for future extensibility.
         """
         if agent.agent_id in self._latest_agent_message:

src/strands/session/session_manager.py

@@ -14,7 +14,7 @@
 from ..types.content import Message
 
 if TYPE_CHECKING:
-    from ..agent.agent import Agent
+    from ..agent.base import AgentBase
     from ..multiagent.base import MultiAgentBase
 
 logger = logging.getLogger(__name__)
@@ -48,40 +48,40 @@ def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None:
         registry.add_callback(AfterMultiAgentInvocationEvent, lambda event: self.sync_multi_agent(event.source))
 
     @abstractmethod
-    def redact_latest_message(self, redact_message: Message, agent: "Agent", **kwargs: Any) -> None:
+    def redact_latest_message(self, redact_message: Message, agent: "AgentBase", **kwargs: Any) -> None:
         """Redact the message most recently appended to the agent in the session.
 
         Args:
             redact_message: New message to use that contains the redact content
-            agent: Agent to apply the message redaction to
+            agent: AgentBase to apply the message redaction to
             **kwargs: Additional keyword arguments for future extensibility.
         """
 
     @abstractmethod
-    def append_message(self, message: Message, agent: "Agent", **kwargs: Any) -> None:
+    def append_message(self, message: Message, agent: "AgentBase", **kwargs: Any) -> None:
         """Append a message to the agent's session.
 
         Args:
             message: Message to add to the agent in the session
-            agent: Agent to append the message to
+            agent: AgentBase to append the message to
             **kwargs: Additional keyword arguments for future extensibility.
         """
 
     @abstractmethod
-    def sync_agent(self, agent: "Agent", **kwargs: Any) -> None:
+    def sync_agent(self, agent: "AgentBase", **kwargs: Any) -> None:
         """Serialize and sync the agent with the session storage.
 
         Args:
-            agent: Agent who should be synchronized with the session storage
+            agent: AgentBase who should be synchronized with the session storage
             **kwargs: Additional keyword arguments for future extensibility.
         """
 
     @abstractmethod
-    def initialize(self, agent: "Agent", **kwargs: Any) -> None:
+    def initialize(self, agent: "AgentBase", **kwargs: Any) -> None:
         """Initialize an agent with a session.
 
         Args:
-            agent: Agent to initialize
+            agent: AgentBase to initialize
             **kwargs: Additional keyword arguments for future extensibility.
         """