Merge branch 'main' into conversation-manager-session

Author: Unshure

src/strands/agent/agent.py

@@ -479,9 +479,10 @@ async def stream_async(self, prompt: Union[str, list[ContentBlock]], **kwargs: A
             prompt: User input as text or list of ContentBlock objects for multi-modal content.
             **kwargs: Additional parameters to pass to the event loop.
 
-        Returns:
+        Yields:
             An async iterator that yields events. Each event is a dictionary containing
             information about the current state of processing, such as:
+
             - data: Text content being generated
             - complete: Whether this is the final chunk
             - current_tool_use: Information about tools being executed

src/strands/event_loop/streaming.py

@@ -253,7 +253,7 @@ async def process_stream(chunks: AsyncIterable[StreamEvent]) -> AsyncGenerator[d
     Args:
         chunks: The chunks of the response stream from the model.
 
-    Returns:
+    Yields:
         The reason for stopping, the constructed message, and the usage metrics.
     """
     stop_reason: StopReason = "end_turn"
@@ -306,7 +306,7 @@ async def stream_messages(
         messages: List of messages to send.
         tool_specs: The list of tool specs.
 
-    Returns:
+    Yields:
         The reason for stopping, the final message, and the usage metrics
     """
     logger.debug("model=<%s> | streaming messages", model)

src/strands/models/model.py

@@ -74,6 +74,7 @@ def stream(
         """Stream conversation with the model.
 
         This method handles the full lifecycle of conversing with the model:
+
         1. Format the messages, tool specs, and configuration into a streaming request
         2. Send the request to the model
         3. Yield the formatted message chunks

src/strands/multiagent/graph.py

@@ -21,7 +21,10 @@
 from dataclasses import dataclass, field
 from typing import Any, Callable, Tuple, cast
 
+from opentelemetry import trace as trace_api
+
 from ..agent import Agent, AgentResult
+from ..telemetry import get_tracer
 from ..types.content import ContentBlock
 from ..types.event_loop import Metrics, Usage
 from .base import MultiAgentBase, MultiAgentResult, NodeResult, Status
@@ -249,6 +252,7 @@ def __init__(self, nodes: dict[str, GraphNode], edges: set[GraphEdge], entry_poi
         self.edges = edges
         self.entry_points = entry_points
         self.state = GraphState()
+        self.tracer = get_tracer()
 
     def execute(self, task: str | list[ContentBlock]) -> GraphResult:
         """Execute task synchronously."""
@@ -274,19 +278,20 @@ async def execute_async(self, task: str | list[ContentBlock]) -> GraphResult:
         )
 
         start_time = time.time()
-        try:
-            await self._execute_graph()
-            self.state.status = Status.COMPLETED
-            logger.debug("status=<%s> | graph execution completed", self.state.status)
-
-        except Exception:
-            logger.exception("graph execution failed")
-            self.state.status = Status.FAILED
-            raise
-        finally:
-            self.state.execution_time = round((time.time() - start_time) * 1000)
-
-        return self._build_result()
+        span = self.tracer.start_multiagent_span(task, "graph")
+        with trace_api.use_span(span, end_on_exit=True):
+            try:
+                await self._execute_graph()
+                self.state.status = Status.COMPLETED
+                logger.debug("status=<%s> | graph execution completed", self.state.status)
+
+            except Exception:
+                logger.exception("graph execution failed")
+                self.state.status = Status.FAILED
+                raise
+            finally:
+                self.state.execution_time = round((time.time() - start_time) * 1000)
+            return self._build_result()
 
     async def _execute_graph(self) -> None:
         """Unified execution flow with conditional routing."""

src/strands/session/file_session_manager.py

@@ -5,7 +5,6 @@
 import os
 import shutil
 import tempfile
-from dataclasses import asdict
 from typing import Any, Optional, cast
 
 from ..types.exceptions import SessionException
@@ -57,22 +56,18 @@ def _get_agent_path(self, session_id: str, agent_id: str) -> str:
         session_path = self._get_session_path(session_id)
         return os.path.join(session_path, "agents", f"{AGENT_PREFIX}{agent_id}")
 
-    def _get_message_path(self, session_id: str, agent_id: str, message_id: str, timestamp: str) -> str:
+    def _get_message_path(self, session_id: str, agent_id: str, message_id: int) -> str:
         """Get message file path.
 
         Args:
             session_id: ID of the session
             agent_id: ID of the agent
-            message_id: ID of the message
-            timestamp: ISO format timestamp to include in filename for sorting
+            message_id: Index of the message
         Returns:
             The filename for the message
         """
         agent_path = self._get_agent_path(session_id, agent_id)
-        # Use timestamp for sortable filenames
-        # Replace colons and periods in ISO format with underscores for filesystem compatibility
-        filename_timestamp = timestamp.replace(":", "_").replace(".", "_")
-        return os.path.join(agent_path, "messages", f"{MESSAGE_PREFIX}{filename_timestamp}_{message_id}.json")
+        return os.path.join(agent_path, "messages", f"{MESSAGE_PREFIX}{message_id}.json")
 
     def _read_file(self, path: str) -> dict[str, Any]:
         """Read JSON file."""
@@ -100,7 +95,7 @@ def create_session(self, session: Session) -> Session:
 
         # Write session file
         session_file = os.path.join(session_dir, "session.json")
-        session_dict = asdict(session)
+        session_dict = session.to_dict()
         self._write_file(session_file, session_dict)
 
         return session
@@ -123,7 +118,7 @@ def create_agent(self, session_id: str, session_agent: SessionAgent) -> None:
         os.makedirs(os.path.join(agent_dir, "messages"), exist_ok=True)
 
         agent_file = os.path.join(agent_dir, "agent.json")
-        session_data = asdict(session_agent)
+        session_data = session_agent.to_dict()
         self._write_file(agent_file, session_data)
 
     def delete_session(self, session_id: str) -> None:
@@ -152,34 +147,25 @@ def update_agent(self, session_id: str, session_agent: SessionAgent) -> None:
 
         session_agent.created_at = previous_agent.created_at
         agent_file = os.path.join(self._get_agent_path(session_id, agent_id), "agent.json")
-        self._write_file(agent_file, asdict(session_agent))
+        self._write_file(agent_file, session_agent.to_dict())
 
     def create_message(self, session_id: str, agent_id: str, session_message: SessionMessage) -> None:
         """Create a new message for the agent."""
         message_file = self._get_message_path(
             session_id,
             agent_id,
             session_message.message_id,
-            session_message.created_at,
         )
-        session_dict = asdict(session_message)
+        session_dict = session_message.to_dict()
         self._write_file(message_file, session_dict)
 
-    def read_message(self, session_id: str, agent_id: str, message_id: str) -> Optional[SessionMessage]:
+    def read_message(self, session_id: str, agent_id: str, message_id: int) -> Optional[SessionMessage]:
         """Read message data."""
-        # Get the messages directory
-        messages_dir = os.path.join(self._get_agent_path(session_id, agent_id), "messages")
-        if not os.path.exists(messages_dir):
+        message_path = self._get_message_path(session_id, agent_id, message_id)
+        if not os.path.exists(message_path):
             return None
-
-        # List files in messages directory, and check if the filename ends with the message id
-        for filename in os.listdir(messages_dir):
-            if filename.endswith(f"{message_id}.json"):
-                file_path = os.path.join(messages_dir, filename)
-                message_data = self._read_file(file_path)
-                return SessionMessage.from_dict(message_data)
-
-        return None
+        message_data = self._read_file(message_path)
+        return SessionMessage.from_dict(message_data)
 
     def update_message(self, session_id: str, agent_id: str, session_message: SessionMessage) -> None:
         """Update message data."""
@@ -190,8 +176,8 @@ def update_message(self, session_id: str, agent_id: str, session_message: Sessio
 
         # Preserve the original created_at timestamp
         session_message.created_at = previous_message.created_at
-        message_file = self._get_message_path(session_id, agent_id, message_id, session_message.created_at)
-        self._write_file(message_file, asdict(session_message))
+        message_file = self._get_message_path(session_id, agent_id, message_id)
+        self._write_file(message_file, session_message.to_dict())
 
     def list_messages(
         self, session_id: str, agent_id: str, limit: Optional[int] = None, offset: int = 0
@@ -201,14 +187,16 @@ def list_messages(
         if not os.path.exists(messages_dir):
             raise SessionException(f"Messages directory missing from agent: {agent_id} in session {session_id}")
 
-        # Read all message files
-        message_files: list[str] = []
+        # Read all message files, and record the index
+        message_index_files: list[tuple[int, str]] = []
         for filename in os.listdir(messages_dir):
             if filename.startswith(MESSAGE_PREFIX) and filename.endswith(".json"):
-                message_files.append(filename)
+                # Extract index from message_<index>.json format
+                index = int(filename[len(MESSAGE_PREFIX) : -5])  # Remove prefix and .json suffix
+                message_index_files.append((index, filename))
 
-        # Sort filenames - the timestamp in the file's name will sort chronologically
-        message_files.sort()
+        # Sort by index and extract just the filenames
+        message_files = [f for _, f in sorted(message_index_files)]
 
         # Apply pagination to filenames
         if limit is not None:

src/strands/session/repository_session_manager.py

@@ -48,11 +48,8 @@ def __init__(
 
         self.session = session
 
-        # Keep track of the initialized agent id's so that two agents in a session cannot share an id
-        self._initialized_agent_ids: set[str] = set()
-
-        # Keep track of the latest message stored in the session in case we need to redact its content.
-        self._latest_message: Optional[SessionMessage] = None
+        # 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) -> None:
         """Append a message to the agent's session.
@@ -61,8 +58,16 @@ def append_message(self, message: Message, agent: Agent) -> None:
             message: Message to add to the agent in the session
             agent: Agent to append the message to
         """
-        self._latest_message = SessionMessage.from_message(message)
-        self.session_repository.create_message(self.session_id, agent.agent_id, self._latest_message)
+        # Calculate the next index (0 if this is the first message, otherwise increment the previous index)
+        latest_agent_message = self._latest_agent_message[agent.agent_id]
+        if latest_agent_message:
+            next_index = latest_agent_message.message_id + 1
+        else:
+            next_index = 0
+
+        session_message = SessionMessage.from_message(message, next_index)
+        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) -> None:
         """Redact the latest message appended to the session.
@@ -71,10 +76,11 @@ def redact_latest_message(self, redact_message: Message, agent: Agent) -> None:
             redact_message: New message to use that contains the redact content
             agent: Agent to apply the message redaction to
         """
-        if self._latest_message is None:
+        latest_agent_message = self._latest_agent_message[agent.agent_id]
+        if latest_agent_message is None:
             raise SessionException("No message to redact.")
-        self._latest_message.redact_message = redact_message
-        return self.session_repository.update_message(self.session_id, agent.agent_id, self._latest_message)
+        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) -> None:
         """Serialize and update the agent into the session repository.
@@ -93,9 +99,9 @@ def initialize(self, agent: Agent) -> None:
         Args:
             agent: Agent to initialize from the session
         """
-        if agent.agent_id in self._initialized_agent_ids:
+        if agent.agent_id in self._latest_agent_message:
             raise SessionException("The `agent_id` of an agent must be unique in a session.")
-        self._initialized_agent_ids.add(agent.agent_id)
+        self._latest_agent_message[agent.agent_id] = None
 
         session_agent = self.session_repository.read_agent(self.session_id, agent.agent_id)
 
@@ -108,8 +114,9 @@ def initialize(self, agent: Agent) -> None:
 
             session_agent = SessionAgent.from_agent(agent)
             self.session_repository.create_agent(self.session_id, session_agent)
-            for message in agent.messages:
-                session_message = SessionMessage.from_message(message)
+            # Initialize messages with sequential indices
+            for i, message in enumerate(agent.messages):
+                session_message = SessionMessage.from_message(message, i)
                 self.session_repository.create_message(self.session_id, agent.agent_id, session_message)
         else:
             logger.debug(