Merge branch 'main' into multiagent-streaming

Author: mkmeral

src/strands/models/openai.py

@@ -214,10 +214,16 @@ def format_request_messages(cls, messages: Messages, system_prompt: Optional[str
         for message in messages:
             contents = message["content"]
 
+            # Check for reasoningContent and warn user
+            if any("reasoningContent" in content for content in contents):
+                logger.warning(
+                    "reasoningContent is not supported in multi-turn conversations with the Chat Completions API."
+                )
+
             formatted_contents = [
                 cls.format_request_message_content(content)
                 for content in contents
-                if not any(block_type in content for block_type in ["toolResult", "toolUse"])
+                if not any(block_type in content for block_type in ["toolResult", "toolUse", "reasoningContent"])
             ]
             formatted_tool_calls = [
                 cls.format_request_message_tool_call(content["toolUse"]) for content in contents if "toolUse" in content
@@ -405,38 +411,46 @@ async def stream(
 
             logger.debug("got response from model")
             yield self.format_chunk({"chunk_type": "message_start"})
-            yield self.format_chunk({"chunk_type": "content_start", "data_type": "text"})
-
             tool_calls: dict[int, list[Any]] = {}
+            data_type = None
+            finish_reason = None  # Store finish_reason for later use
+            event = None  # Initialize for scope safety
 
             async for event in response:
                 # Defensive: skip events with empty or missing choices
                 if not getattr(event, "choices", None):
                     continue
                 choice = event.choices[0]
 
-                if choice.delta.content:
-                    yield self.format_chunk(
-                        {"chunk_type": "content_delta", "data_type": "text", "data": choice.delta.content}
-                    )
-
                 if hasattr(choice.delta, "reasoning_content") and choice.delta.reasoning_content:
+                    chunks, data_type = self._stream_switch_content("reasoning_content", data_type)
+                    for chunk in chunks:
+                        yield chunk
                     yield self.format_chunk(
                         {
                             "chunk_type": "content_delta",
-                            "data_type": "reasoning_content",
+                            "data_type": data_type,
                             "data": choice.delta.reasoning_content,
                         }
                     )
 
+                if choice.delta.content:
+                    chunks, data_type = self._stream_switch_content("text", data_type)
+                    for chunk in chunks:
+                        yield chunk
+                    yield self.format_chunk(
+                        {"chunk_type": "content_delta", "data_type": data_type, "data": choice.delta.content}
+                    )
+
                 for tool_call in choice.delta.tool_calls or []:
                     tool_calls.setdefault(tool_call.index, []).append(tool_call)
 
                 if choice.finish_reason:
+                    finish_reason = choice.finish_reason  # Store for use outside loop
+                    if data_type:
+                        yield self.format_chunk({"chunk_type": "content_stop", "data_type": data_type})
                     break
 
-            yield self.format_chunk({"chunk_type": "content_stop", "data_type": "text"})
-
             for tool_deltas in tool_calls.values():
                 yield self.format_chunk({"chunk_type": "content_start", "data_type": "tool", "data": tool_deltas[0]})
 
@@ -445,17 +459,37 @@ async def stream(
 
                 yield self.format_chunk({"chunk_type": "content_stop", "data_type": "tool"})
 
-            yield self.format_chunk({"chunk_type": "message_stop", "data": choice.finish_reason})
+            yield self.format_chunk({"chunk_type": "message_stop", "data": finish_reason or "end_turn"})
 
             # Skip remaining events as we don't have use for anything except the final usage payload
             async for event in response:
                 _ = event
 
-            if event.usage:
+            if event and hasattr(event, "usage") and event.usage:
                 yield self.format_chunk({"chunk_type": "metadata", "data": event.usage})
 
         logger.debug("finished streaming response from model")
 
+    def _stream_switch_content(self, data_type: str, prev_data_type: str | None) -> tuple[list[StreamEvent], str]:
+        """Handle switching to a new content stream.
+
+        Args:
+            data_type: The next content data type.
+            prev_data_type: The previous content data type.
+
+        Returns:
+            Tuple containing:
+            - Stop block for previous content and the start block for the next content.
+            - Next content data type.
+        """
+        chunks = []
+        if data_type != prev_data_type:
+            if prev_data_type is not None:
+                chunks.append(self.format_chunk({"chunk_type": "content_stop", "data_type": prev_data_type}))
+            chunks.append(self.format_chunk({"chunk_type": "content_start", "data_type": data_type}))
+
+        return chunks, data_type
+
     @override
     async def structured_output(
         self, output_model: Type[T], prompt: Messages, system_prompt: Optional[str] = None, **kwargs: Any

src/strands/multiagent/base.py

@@ -137,7 +137,7 @@ def from_dict(cls, data: dict[str, Any]) -> "MultiAgentResult":
         metrics = _parse_metrics(data.get("accumulated_metrics", {}))
 
         multiagent_result = cls(
-            status=Status(data.get("status", Status.PENDING.value)),
+            status=Status(data["status"]),
             results=results,
             accumulated_usage=usage,
             accumulated_metrics=metrics,
@@ -164,8 +164,13 @@ class MultiAgentBase(ABC):
 
     This class integrates with existing Strands Agent instances and provides
     multi-agent orchestration capabilities.
+
+    Attributes:
+        id: Unique MultiAgent id for session management,etc.
     """
 
+    id: str
+
     @abstractmethod
     async def invoke_async(
         self, task: str | list[ContentBlock], invocation_state: dict[str, Any] | None = None, **kwargs: Any

src/strands/multiagent/swarm.py

@@ -133,7 +133,7 @@ def _validate_json_serializable(self, value: Any) -> None:
 class SwarmState:
     """Current state of swarm execution."""
 
-    current_node: SwarmNode  # The agent currently executing
+    current_node: SwarmNode | None  # The agent currently executing
     task: str | list[ContentBlock]  # The original task from the user that is being executed
     completion_status: Status = Status.PENDING  # Current swarm execution status
     shared_context: SharedContext = field(default_factory=SharedContext)  # Context shared between agents
@@ -238,7 +238,7 @@ def __init__(
         self.shared_context = SharedContext()
         self.nodes: dict[str, SwarmNode] = {}
         self.state = SwarmState(
-            current_node=SwarmNode("", Agent()),  # Placeholder, will be set properly
+            current_node=None,  # Placeholder, will be set properly
             task="",
             completion_status=Status.PENDING,
         )
@@ -328,7 +328,8 @@ async def stream_async(
         span = self.tracer.start_multiagent_span(task, "swarm")
         with trace_api.use_span(span, end_on_exit=True):
             try:
-                logger.debug("current_node=<%s> | starting swarm execution with node", self.state.current_node.node_id)
+                current_node = cast(SwarmNode, self.state.current_node)
+                logger.debug("current_node=<%s> | starting swarm execution with node", current_node.node_id)
                 logger.debug(
                     "max_handoffs=<%d>, max_iterations=<%d>, timeout=<%s>s | swarm execution config",
                     self.max_handoffs,
@@ -522,7 +523,7 @@ def _handle_handoff(self, target_node: SwarmNode, message: str, context: dict[st
             return
 
         # Update swarm state
-        previous_agent = self.state.current_node
+        previous_agent = cast(SwarmNode, self.state.current_node)
         self.state.current_node = target_node
 
         # Store handoff message for the target agent

src/strands/session/file_session_manager.py

@@ -5,19 +5,23 @@
 import os
 import shutil
 import tempfile
-from typing import Any, Optional, cast
+from typing import TYPE_CHECKING, Any, Optional, cast
 
 from .. import _identifier
 from ..types.exceptions import SessionException
 from ..types.session import Session, SessionAgent, SessionMessage
 from .repository_session_manager import RepositorySessionManager
 from .session_repository import SessionRepository
 
+if TYPE_CHECKING:
+    from ..multiagent.base import MultiAgentBase
+
 logger = logging.getLogger(__name__)
 
 SESSION_PREFIX = "session_"
 AGENT_PREFIX = "agent_"
 MESSAGE_PREFIX = "message_"
+MULTI_AGENT_PREFIX = "multi_agent_"
 
 
 class FileSessionManager(RepositorySessionManager, SessionRepository):
@@ -37,7 +41,12 @@ class FileSessionManager(RepositorySessionManager, SessionRepository):
     ```
     """
 
-    def __init__(self, session_id: str, storage_dir: Optional[str] = None, **kwargs: Any):
+    def __init__(
+        self,
+        session_id: str,
+        storage_dir: Optional[str] = None,
+        **kwargs: Any,
+    ):
         """Initialize FileSession with filesystem storage.
 
         Args:
@@ -107,8 +116,11 @@ def _read_file(self, path: str) -> dict[str, Any]:
     def _write_file(self, path: str, data: dict[str, Any]) -> None:
         """Write JSON file."""
         os.makedirs(os.path.dirname(path), exist_ok=True)
-        with open(path, "w", encoding="utf-8") as f:
+        # This automic write ensure the completeness of session files in both single agent/ multi agents
+        tmp = f"{path}.tmp"
+        with open(tmp, "w", encoding="utf-8", newline="\n") as f:
             json.dump(data, f, indent=2, ensure_ascii=False)
+        os.replace(tmp, path)
 
     def create_session(self, session: Session, **kwargs: Any) -> Session:
         """Create a new session."""
@@ -119,6 +131,7 @@ def create_session(self, session: Session, **kwargs: Any) -> Session:
         # Create directory structure
         os.makedirs(session_dir, exist_ok=True)
         os.makedirs(os.path.join(session_dir, "agents"), exist_ok=True)
+        os.makedirs(os.path.join(session_dir, "multi_agents"), exist_ok=True)
 
         # Write session file
         session_file = os.path.join(session_dir, "session.json")
@@ -239,3 +252,36 @@ def list_messages(
             messages.append(SessionMessage.from_dict(message_data))
 
         return messages
+
+    def _get_multi_agent_path(self, session_id: str, multi_agent_id: str) -> str:
+        """Get multi-agent state file path."""
+        session_path = self._get_session_path(session_id)
+        multi_agent_id = _identifier.validate(multi_agent_id, _identifier.Identifier.AGENT)
+        return os.path.join(session_path, "multi_agents", f"{MULTI_AGENT_PREFIX}{multi_agent_id}")
+
+    def create_multi_agent(self, session_id: str, multi_agent: "MultiAgentBase", **kwargs: Any) -> None:
+        """Create a new multiagent state in the session."""
+        multi_agent_id = multi_agent.id
+        multi_agent_dir = self._get_multi_agent_path(session_id, multi_agent_id)
+        os.makedirs(multi_agent_dir, exist_ok=True)
+
+        multi_agent_file = os.path.join(multi_agent_dir, "multi_agent.json")
+        session_data = multi_agent.serialize_state()
+        self._write_file(multi_agent_file, session_data)
+
+    def read_multi_agent(self, session_id: str, multi_agent_id: str, **kwargs: Any) -> Optional[dict[str, Any]]:
+        """Read multi-agent state from filesystem."""
+        multi_agent_file = os.path.join(self._get_multi_agent_path(session_id, multi_agent_id), "multi_agent.json")
+        if not os.path.exists(multi_agent_file):
+            return None
+        return self._read_file(multi_agent_file)
+
+    def update_multi_agent(self, session_id: str, multi_agent: "MultiAgentBase", **kwargs: Any) -> None:
+        """Update multi-agent state from filesystem."""
+        multi_agent_state = multi_agent.serialize_state()
+        previous_multi_agent_state = self.read_multi_agent(session_id=session_id, multi_agent_id=multi_agent.id)
+        if previous_multi_agent_state is None:
+            raise SessionException(f"MultiAgent state {multi_agent.id} in session {session_id} does not exist")
+
+        multi_agent_file = os.path.join(self._get_multi_agent_path(session_id, multi_agent.id), "multi_agent.json")
+        self._write_file(multi_agent_file, multi_agent_state)

src/strands/session/repository_session_manager.py

@@ -17,14 +17,20 @@
 
 if TYPE_CHECKING:
     from ..agent.agent import Agent
+    from ..multiagent.base import MultiAgentBase
 
 logger = logging.getLogger(__name__)
 
 
 class RepositorySessionManager(SessionManager):
     """Session manager for persisting agents in a SessionRepository."""
 
-    def __init__(self, session_id: str, session_repository: SessionRepository, **kwargs: Any):
+    def __init__(
+        self,
+        session_id: str,
+        session_repository: SessionRepository,
+        **kwargs: Any,
+    ):
         """Initialize the RepositorySessionManager.
 
         If no session with the specified session_id exists yet, it will be created
@@ -152,3 +158,26 @@ def initialize(self, agent: "Agent", **kwargs: Any) -> None:
 
             # Restore the agents messages array including the optional prepend messages
             agent.messages = prepend_messages + [session_message.to_message() for session_message in session_messages]
+
+    def sync_multi_agent(self, source: "MultiAgentBase", **kwargs: Any) -> None:
+        """Serialize and update the multi-agent state into the session repository.
+
+        Args:
+            source: Multi-agent source object to sync to the session.
+            **kwargs: Additional keyword arguments for future extensibility.
+        """
+        self.session_repository.update_multi_agent(self.session_id, source)
+
+    def initialize_multi_agent(self, source: "MultiAgentBase", **kwargs: Any) -> None:
+        """Initialize multi-agent state from the session repository.
+
+        Args:
+            source: Multi-agent source object to restore state into
+            **kwargs: Additional keyword arguments for future extensibility.
+        """
+        state = self.session_repository.read_multi_agent(self.session_id, source.id, **kwargs)
+        if state is None:
+            self.session_repository.create_multi_agent(self.session_id, source, **kwargs)
+        else:
+            logger.debug("session_id=<%s> | restoring multi-agent state", self.session_id)
+            source.deserialize_state(state)

src/strands/session/s3_session_manager.py

@@ -2,7 +2,7 @@
 
 import json
 import logging
-from typing import Any, Dict, List, Optional, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
 
 import boto3
 from botocore.config import Config as BotocoreConfig
@@ -14,11 +14,15 @@
 from .repository_session_manager import RepositorySessionManager
 from .session_repository import SessionRepository
 
+if TYPE_CHECKING:
+    from ..multiagent.base import MultiAgentBase
+
 logger = logging.getLogger(__name__)
 
 SESSION_PREFIX = "session_"
 AGENT_PREFIX = "agent_"
 MESSAGE_PREFIX = "message_"
+MULTI_AGENT_PREFIX = "multi_agent_"
 
 
 class S3SessionManager(RepositorySessionManager, SessionRepository):
@@ -294,3 +298,31 @@ def list_messages(
 
         except ClientError as e:
             raise SessionException(f"S3 error reading messages: {e}") from e
+
+    def _get_multi_agent_path(self, session_id: str, multi_agent_id: str) -> str:
+        """Get multi-agent S3 prefix."""
+        session_path = self._get_session_path(session_id)
+        multi_agent_id = _identifier.validate(multi_agent_id, _identifier.Identifier.AGENT)
+        return f"{session_path}multi_agents/{MULTI_AGENT_PREFIX}{multi_agent_id}/"
+
+    def create_multi_agent(self, session_id: str, multi_agent: "MultiAgentBase", **kwargs: Any) -> None:
+        """Create a new multiagent state in S3."""
+        multi_agent_id = multi_agent.id
+        multi_agent_key = f"{self._get_multi_agent_path(session_id, multi_agent_id)}multi_agent.json"
+        session_data = multi_agent.serialize_state()
+        self._write_s3_object(multi_agent_key, session_data)
+
+    def read_multi_agent(self, session_id: str, multi_agent_id: str, **kwargs: Any) -> Optional[dict[str, Any]]:
+        """Read multi-agent state from S3."""
+        multi_agent_key = f"{self._get_multi_agent_path(session_id, multi_agent_id)}multi_agent.json"
+        return self._read_s3_object(multi_agent_key)
+
+    def update_multi_agent(self, session_id: str, multi_agent: "MultiAgentBase", **kwargs: Any) -> None:
+        """Update multi-agent state in S3."""
+        multi_agent_state = multi_agent.serialize_state()
+        previous_multi_agent_state = self.read_multi_agent(session_id=session_id, multi_agent_id=multi_agent.id)
+        if previous_multi_agent_state is None:
+            raise SessionException(f"MultiAgent state {multi_agent.id} in session {session_id} does not exist")
+
+        multi_agent_key = f"{self._get_multi_agent_path(session_id, multi_agent.id)}multi_agent.json"
+        self._write_s3_object(multi_agent_key, multi_agent_state)