fix

Author: awsarron

src/strands/agent/base.py

@@ -7,7 +7,6 @@
 
 from ..types.agent import AgentInput
 from .agent_result import AgentResult
-from .state import AgentState
 
 
 @runtime_checkable
@@ -18,15 +17,6 @@ class AgentBase(Protocol):
     must satisfy.
     """
 
-    agent_id: str
-    """Unique identifier for the agent."""
-
-    name: str
-    """Human-readable name of the agent."""
-
-    state: AgentState
-    """Current state of the agent."""
-
     async def invoke_async(
         self,
         prompt: AgentInput = None,

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, AgentBase
+from ..agent import Agent
 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: AgentBase
+    executor: Agent
     _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[AgentBase],
+        nodes: list[Agent],
         *,
-        entry_point: AgentBase | None = None,
+        entry_point: Agent | 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. AgentBase) to include in the swarm
-            entry_point: AgentBase to start with. If None, uses the first agent (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)
             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[AgentBase]) -> None:
+    def _setup_swarm(self, nodes: list[Agent]) -> None:
         """Initialize swarm configuration."""
         # Validate nodes before setup
         self._validate_swarm(nodes)
@@ -467,7 +467,7 @@ def _setup_swarm(self, nodes: list[AgentBase]) -> 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[AgentBase]) -> None:
+    def _validate_swarm(self, nodes: list[Agent]) -> 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.base import AgentBase
+    from ..agent.agent import Agent
     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: "AgentBase", **kwargs: Any) -> None:
+    def append_message(self, message: Message, agent: "Agent", **kwargs: Any) -> None:
         """Append a message to the agent's session.
 
         Args:
             message: Message to add to the agent in the session
-            agent: AgentBase to append the message to
+            agent: Agent 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: "AgentBase", **kwargs: Any) ->
         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: "AgentBase", **kwargs: Any) -> None:
+    def redact_latest_message(self, redact_message: Message, agent: "Agent", **kwargs: Any) -> None:
         """Redact the latest message appended to the session.
 
         Args:
             redact_message: New message to use that contains the redact content
-            agent: AgentBase to apply the message redaction to
+            agent: Agent 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: "AgentBase", **k
         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: "AgentBase", **kwargs: Any) -> None:
+    def sync_agent(self, agent: "Agent", **kwargs: Any) -> None:
         """Serialize and update the agent into the session repository.
 
         Args:
-            agent: AgentBase to sync to the session.
+            agent: Agent 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: "AgentBase", **kwargs: Any) -> None:
+    def initialize(self, agent: "Agent", **kwargs: Any) -> None:
         """Initialize an agent with a session.
 
         Args:
-            agent: AgentBase to initialize from the session
+            agent: Agent 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.base import AgentBase
+    from ..agent.agent import Agent
     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: "AgentBase", **kwargs: Any) -> None:
+    def redact_latest_message(self, redact_message: Message, agent: "Agent", **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: AgentBase to apply the message redaction to
+            agent: Agent to apply the message redaction to
             **kwargs: Additional keyword arguments for future extensibility.
         """
 
     @abstractmethod
-    def append_message(self, message: Message, agent: "AgentBase", **kwargs: Any) -> None:
+    def append_message(self, message: Message, agent: "Agent", **kwargs: Any) -> None:
         """Append a message to the agent's session.
 
         Args:
             message: Message to add to the agent in the session
-            agent: AgentBase to append the message to
+            agent: Agent to append the message to
             **kwargs: Additional keyword arguments for future extensibility.
         """
 
     @abstractmethod
-    def sync_agent(self, agent: "AgentBase", **kwargs: Any) -> None:
+    def sync_agent(self, agent: "Agent", **kwargs: Any) -> None:
         """Serialize and sync the agent with the session storage.
 
         Args:
-            agent: AgentBase who should be synchronized with the session storage
+            agent: Agent who should be synchronized with the session storage
             **kwargs: Additional keyword arguments for future extensibility.
         """
 
     @abstractmethod
-    def initialize(self, agent: "AgentBase", **kwargs: Any) -> None:
+    def initialize(self, agent: "Agent", **kwargs: Any) -> None:
         """Initialize an agent with a session.
 
         Args:
-            agent: AgentBase to initialize
+            agent: Agent to initialize
             **kwargs: Additional keyword arguments for future extensibility.
         """