agentrun-sdk 0.1.2__py3-none-any.whl

agentrun_sdk/agent/agent_result.py

@@ -0,0 +1,46 @@
+"""Agent result handling for SDK.
+
+This module defines the AgentResult class which encapsulates the complete response from an agent's processing cycle.
+"""
+
+from dataclasses import dataclass
+from typing import Any
+
+from ..telemetry.metrics import EventLoopMetrics
+from ..types.content import Message
+from ..types.streaming import StopReason
+
+
+@dataclass
+class AgentResult:
+    """Represents the last result of invoking an agent with a prompt.
+
+    Attributes:
+        stop_reason: The reason why the agent's processing stopped.
+        message: The last message generated by the agent.
+        metrics: Performance metrics collected during processing.
+        state: Additional state information from the event loop.
+    """
+
+    stop_reason: StopReason
+    message: Message
+    metrics: EventLoopMetrics
+    state: Any
+
+    def __str__(self) -> str:
+        """Get the agent's last message as a string.
+
+        This method extracts and concatenates all text content from the final message, ignoring any non-text content
+        like images or structured data.
+
+        Returns:
+            The agent's last message as a string.
+        """
+        content_array = self.message.get("content", [])
+
+        result = ""
+        for item in content_array:
+            if isinstance(item, dict) and "text" in item:
+                result += item.get("text", "") + "\n"
+
+        return result

agentrun_sdk/agent/conversation_manager/__init__.py

@@ -0,0 +1,26 @@
+"""This package provides classes for managing conversation history during agent execution.
+
+It includes:
+
+- ConversationManager: Abstract base class defining the conversation management interface
+- NullConversationManager: A no-op implementation that does not modify conversation history
+- SlidingWindowConversationManager: An implementation that maintains a sliding window of messages to control context
+  size while preserving conversation coherence
+- SummarizingConversationManager: An implementation that summarizes older context instead
+  of simply trimming it
+
+Conversation managers help control memory usage and context length while maintaining relevant conversation state, which
+is critical for effective agent interactions.
+"""
+
+from .conversation_manager import ConversationManager
+from .null_conversation_manager import NullConversationManager
+from .sliding_window_conversation_manager import SlidingWindowConversationManager
+from .summarizing_conversation_manager import SummarizingConversationManager
+
+__all__ = [
+    "ConversationManager",
+    "NullConversationManager",
+    "SlidingWindowConversationManager",
+    "SummarizingConversationManager",
+]

agentrun_sdk/agent/conversation_manager/conversation_manager.py

@@ -0,0 +1,88 @@
+"""Abstract interface for conversation history management."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Optional
+
+from ...types.content import Message
+
+if TYPE_CHECKING:
+    from ...agent.agent import Agent
+
+
+class ConversationManager(ABC):
+    """Abstract base class for managing conversation history.
+
+    This class provides an interface for implementing conversation management strategies to control the size of message
+    arrays/conversation histories, helping to:
+
+    - Manage memory usage
+    - Control context length
+    - Maintain relevant conversation state
+    """
+
+    def __init__(self) -> None:
+        """Initialize the ConversationManager.
+
+        Attributes:
+          removed_message_count: The messages that have been removed from the agents messages array.
+              These represent messages provided by the user or LLM that have been removed, not messages
+              included by the conversation manager through something like summarization.
+        """
+        self.removed_message_count = 0
+
+    def restore_from_session(self, state: dict[str, Any]) -> Optional[list[Message]]:
+        """Restore the Conversation Manager's state from a session.
+
+        Args:
+            state: Previous state of the conversation manager
+        Returns:
+            Optional list of messages to prepend to the agents messages. By default returns None.
+        """
+        if state.get("__name__") != self.__class__.__name__:
+            raise ValueError("Invalid conversation manager state.")
+        self.removed_message_count = state["removed_message_count"]
+        return None
+
+    def get_state(self) -> dict[str, Any]:
+        """Get the current state of a Conversation Manager as a Json serializable dictionary."""
+        return {
+            "__name__": self.__class__.__name__,
+            "removed_message_count": self.removed_message_count,
+        }
+
+    @abstractmethod
+    def apply_management(self, agent: "Agent", **kwargs: Any) -> None:
+        """Applies management strategy to the provided agent.
+
+        Processes the conversation history to maintain appropriate size by modifying the messages list in-place.
+        Implementations should handle message pruning, summarization, or other size management techniques to keep the
+        conversation context within desired bounds.
+
+        Args:
+            agent: The agent whose conversation history will be manage.
+                This list is modified in-place.
+            **kwargs: Additional keyword arguments for future extensibility.
+        """
+        pass
+
+    @abstractmethod
+    def reduce_context(self, agent: "Agent", e: Optional[Exception] = None, **kwargs: Any) -> None:
+        """Called when the model's context window is exceeded.
+
+        This method should implement the specific strategy for reducing the window size when a context overflow occurs.
+        It is typically called after a ContextWindowOverflowException is caught.
+
+        Implementations might use strategies such as:
+
+        - Removing the N oldest messages
+        - Summarizing older context
+        - Applying importance-based filtering
+        - Maintaining critical conversation markers
+
+        Args:
+            agent: The agent whose conversation history will be reduced.
+                This list is modified in-place.
+            e: The exception that triggered the context reduction, if any.
+            **kwargs: Additional keyword arguments for future extensibility.
+        """
+        pass

agentrun_sdk/agent/conversation_manager/null_conversation_manager.py

@@ -0,0 +1,46 @@
+"""Null implementation of conversation management."""
+
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    from ...agent.agent import Agent
+
+from ...types.exceptions import ContextWindowOverflowException
+from .conversation_manager import ConversationManager
+
+
+class NullConversationManager(ConversationManager):
+    """A no-op conversation manager that does not modify the conversation history.
+
+    Useful for:
+
+    - Testing scenarios where conversation management should be disabled
+    - Cases where conversation history is managed externally
+    - Situations where the full conversation history should be preserved
+    """
+
+    def apply_management(self, agent: "Agent", **kwargs: Any) -> None:
+        """Does nothing to the conversation history.
+
+        Args:
+            agent: The agent whose conversation history will remain unmodified.
+            **kwargs: Additional keyword arguments for future extensibility.
+        """
+        pass
+
+    def reduce_context(self, agent: "Agent", e: Optional[Exception] = None, **kwargs: Any) -> None:
+        """Does not reduce context and raises an exception.
+
+        Args:
+            agent: The agent whose conversation history will remain unmodified.
+            e: The exception that triggered the context reduction, if any.
+            **kwargs: Additional keyword arguments for future extensibility.
+
+        Raises:
+            e: If provided.
+            ContextWindowOverflowException: If e is None.
+        """
+        if e:
+            raise e
+        else:
+            raise ContextWindowOverflowException("Context window overflowed!")

agentrun_sdk/agent/conversation_manager/sliding_window_conversation_manager.py

@@ -0,0 +1,179 @@
+"""Sliding window conversation history management."""
+
+import logging
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    from ...agent.agent import Agent
+
+from ...types.content import Messages
+from ...types.exceptions import ContextWindowOverflowException
+from .conversation_manager import ConversationManager
+
+logger = logging.getLogger(__name__)
+
+
+class SlidingWindowConversationManager(ConversationManager):
+    """Implements a sliding window strategy for managing conversation history.
+
+    This class handles the logic of maintaining a conversation window that preserves tool usage pairs and avoids
+    invalid window states.
+    """
+
+    def __init__(self, window_size: int = 40, should_truncate_results: bool = True):
+        """Initialize the sliding window conversation manager.
+
+        Args:
+            window_size: Maximum number of messages to keep in the agent's history.
+                Defaults to 40 messages.
+            should_truncate_results: Truncate tool results when a message is too large for the model's context window
+        """
+        super().__init__()
+        self.window_size = window_size
+        self.should_truncate_results = should_truncate_results
+
+    def apply_management(self, agent: "Agent", **kwargs: Any) -> None:
+        """Apply the sliding window to the agent's messages array to maintain a manageable history size.
+
+        This method is called after every event loop cycle to apply a sliding window if the message count
+        exceeds the window size.
+
+        Args:
+            agent: The agent whose messages will be managed.
+                This list is modified in-place.
+            **kwargs: Additional keyword arguments for future extensibility.
+        """
+        messages = agent.messages
+
+        if len(messages) <= self.window_size:
+            logger.debug(
+                "message_count=<%s>, window_size=<%s> | skipping context reduction", len(messages), self.window_size
+            )
+            return
+        self.reduce_context(agent)
+
+    def reduce_context(self, agent: "Agent", e: Optional[Exception] = None, **kwargs: Any) -> None:
+        """Trim the oldest messages to reduce the conversation context size.
+
+        The method handles special cases where trimming the messages leads to:
+         - toolResult with no corresponding toolUse
+         - toolUse with no corresponding toolResult
+
+        Args:
+            agent: The agent whose messages will be reduce.
+                This list is modified in-place.
+            e: The exception that triggered the context reduction, if any.
+            **kwargs: Additional keyword arguments for future extensibility.
+
+        Raises:
+            ContextWindowOverflowException: If the context cannot be reduced further.
+                Such as when the conversation is already minimal or when tool result messages cannot be properly
+                converted.
+        """
+        messages = agent.messages
+
+        # Try to truncate the tool result first
+        last_message_idx_with_tool_results = self._find_last_message_with_tool_results(messages)
+        if last_message_idx_with_tool_results is not None and self.should_truncate_results:
+            logger.debug(
+                "message_index=<%s> | found message with tool results at index", last_message_idx_with_tool_results
+            )
+            results_truncated = self._truncate_tool_results(messages, last_message_idx_with_tool_results)
+            if results_truncated:
+                logger.debug("message_index=<%s> | tool results truncated", last_message_idx_with_tool_results)
+                return
+
+        # Try to trim index id when tool result cannot be truncated anymore
+        # If the number of messages is less than the window_size, then we default to 2, otherwise, trim to window size
+        trim_index = 2 if len(messages) <= self.window_size else len(messages) - self.window_size
+
+        # Find the next valid trim_index
+        while trim_index < len(messages):
+            if (
+                # Oldest message cannot be a toolResult because it needs a toolUse preceding it
+                any("toolResult" in content for content in messages[trim_index]["content"])
+                or (
+                    # Oldest message can be a toolUse only if a toolResult immediately follows it.
+                    any("toolUse" in content for content in messages[trim_index]["content"])
+                    and trim_index + 1 < len(messages)
+                    and not any("toolResult" in content for content in messages[trim_index + 1]["content"])
+                )
+            ):
+                trim_index += 1
+            else:
+                break
+        else:
+            # If we didn't find a valid trim_index, then we throw
+            raise ContextWindowOverflowException("Unable to trim conversation context!") from e
+
+        # trim_index represents the number of messages being removed from the agents messages array
+        self.removed_message_count += trim_index
+
+        # Overwrite message history
+        messages[:] = messages[trim_index:]
+
+    def _truncate_tool_results(self, messages: Messages, msg_idx: int) -> bool:
+        """Truncate tool results in a message to reduce context size.
+
+        When a message contains tool results that are too large for the model's context window, this function
+        replaces the content of those tool results with a simple error message.
+
+        Args:
+            messages: The conversation message history.
+            msg_idx: Index of the message containing tool results to truncate.
+
+        Returns:
+            True if any changes were made to the message, False otherwise.
+        """
+        if msg_idx >= len(messages) or msg_idx < 0:
+            return False
+
+        message = messages[msg_idx]
+        changes_made = False
+        tool_result_too_large_message = "The tool result was too large!"
+        for i, content in enumerate(message.get("content", [])):
+            if isinstance(content, dict) and "toolResult" in content:
+                tool_result_content_text = next(
+                    (item["text"] for item in content["toolResult"]["content"] if "text" in item),
+                    "",
+                )
+                # make the overwriting logic togglable
+                if (
+                    message["content"][i]["toolResult"]["status"] == "error"
+                    and tool_result_content_text == tool_result_too_large_message
+                ):
+                    logger.info("ToolResult has already been updated, skipping overwrite")
+                    return False
+                # Update status to error with informative message
+                message["content"][i]["toolResult"]["status"] = "error"
+                message["content"][i]["toolResult"]["content"] = [{"text": tool_result_too_large_message}]
+                changes_made = True
+
+        return changes_made
+
+    def _find_last_message_with_tool_results(self, messages: Messages) -> Optional[int]:
+        """Find the index of the last message containing tool results.
+
+        This is useful for identifying messages that might need to be truncated to reduce context size.
+
+        Args:
+            messages: The conversation message history.
+
+        Returns:
+            Index of the last message with tool results, or None if no such message exists.
+        """
+        # Iterate backwards through all messages (from newest to oldest)
+        for idx in range(len(messages) - 1, -1, -1):
+            # Check if this message has any content with toolResult
+            current_message = messages[idx]
+            has_tool_result = False
+
+            for content in current_message.get("content", []):
+                if isinstance(content, dict) and "toolResult" in content:
+                    has_tool_result = True
+                    break
+
+            if has_tool_result:
+                return idx
+
+        return None

agentrun_sdk/agent/conversation_manager/summarizing_conversation_manager.py

@@ -0,0 +1,252 @@
+"""Summarizing conversation history management with configurable options."""
+
+import logging
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from typing_extensions import override
+
+from ...types.content import Message
+from ...types.exceptions import ContextWindowOverflowException
+from .conversation_manager import ConversationManager
+
+if TYPE_CHECKING:
+    from ..agent import Agent
+
+
+logger = logging.getLogger(__name__)
+
+
+DEFAULT_SUMMARIZATION_PROMPT = """You are a conversation summarizer. Provide a concise summary of the conversation \
+history.
+
+Format Requirements:
+- You MUST create a structured and concise summary in bullet-point format.
+- You MUST NOT respond conversationally.
+- You MUST NOT address the user directly.
+
+Task:
+Your task is to create a structured summary document:
+- It MUST contain bullet points with key topics and questions covered
+- It MUST contain bullet points for all significant tools executed and their results
+- It MUST contain bullet points for any code or technical information shared
+- It MUST contain a section of key insights gained
+- It MUST format the summary in the third person
+
+Example format:
+
+## Conversation Summary
+* Topic 1: Key information
+* Topic 2: Key information
+*
+## Tools Executed
+* Tool X: Result Y"""
+
+
+class SummarizingConversationManager(ConversationManager):
+    """Implements a summarizing window manager.
+
+    This manager provides a configurable option to summarize older context instead of
+    simply trimming it, helping preserve important information while staying within
+    context limits.
+    """
+
+    def __init__(
+        self,
+        summary_ratio: float = 0.3,
+        preserve_recent_messages: int = 10,
+        summarization_agent: Optional["Agent"] = None,
+        summarization_system_prompt: Optional[str] = None,
+    ):
+        """Initialize the summarizing conversation manager.
+
+        Args:
+            summary_ratio: Ratio of messages to summarize vs keep when context overflow occurs.
+                Value between 0.1 and 0.8. Defaults to 0.3 (summarize 30% of oldest messages).
+            preserve_recent_messages: Minimum number of recent messages to always keep.
+                Defaults to 10 messages.
+            summarization_agent: Optional agent to use for summarization instead of the parent agent.
+                If provided, this agent can use tools as part of the summarization process.
+            summarization_system_prompt: Optional system prompt override for summarization.
+                If None, uses the default summarization prompt.
+        """
+        super().__init__()
+        if summarization_agent is not None and summarization_system_prompt is not None:
+            raise ValueError(
+                "Cannot provide both summarization_agent and summarization_system_prompt. "
+                "Agents come with their own system prompt."
+            )
+
+        self.summary_ratio = max(0.1, min(0.8, summary_ratio))
+        self.preserve_recent_messages = preserve_recent_messages
+        self.summarization_agent = summarization_agent
+        self.summarization_system_prompt = summarization_system_prompt
+        self._summary_message: Optional[Message] = None
+
+    @override
+    def restore_from_session(self, state: dict[str, Any]) -> Optional[list[Message]]:
+        """Restores the Summarizing Conversation manager from its previous state in a session.
+
+        Args:
+            state: The previous state of the Summarizing Conversation Manager.
+
+        Returns:
+            Optionally returns the previous conversation summary if it exists.
+        """
+        super().restore_from_session(state)
+        self._summary_message = state.get("summary_message")
+        return [self._summary_message] if self._summary_message else None
+
+    def get_state(self) -> dict[str, Any]:
+        """Returns a dictionary representation of the state for the Summarizing Conversation Manager."""
+        return {"summary_message": self._summary_message, **super().get_state()}
+
+    def apply_management(self, agent: "Agent", **kwargs: Any) -> None:
+        """Apply management strategy to conversation history.
+
+        For the summarizing conversation manager, no proactive management is performed.
+        Summarization only occurs when there's a context overflow that triggers reduce_context.
+
+        Args:
+            agent: The agent whose conversation history will be managed.
+                The agent's messages list is modified in-place.
+            **kwargs: Additional keyword arguments for future extensibility.
+        """
+        # No proactive management - summarization only happens on context overflow
+        pass
+
+    def reduce_context(self, agent: "Agent", e: Optional[Exception] = None, **kwargs: Any) -> None:
+        """Reduce context using summarization.
+
+        Args:
+            agent: The agent whose conversation history will be reduced.
+                The agent's messages list is modified in-place.
+            e: The exception that triggered the context reduction, if any.
+            **kwargs: Additional keyword arguments for future extensibility.
+
+        Raises:
+            ContextWindowOverflowException: If the context cannot be summarized.
+        """
+        try:
+            # Calculate how many messages to summarize
+            messages_to_summarize_count = max(1, int(len(agent.messages) * self.summary_ratio))
+
+            # Ensure we don't summarize recent messages
+            messages_to_summarize_count = min(
+                messages_to_summarize_count, len(agent.messages) - self.preserve_recent_messages
+            )
+
+            if messages_to_summarize_count <= 0:
+                raise ContextWindowOverflowException("Cannot summarize: insufficient messages for summarization")
+
+            # Adjust split point to avoid breaking ToolUse/ToolResult pairs
+            messages_to_summarize_count = self._adjust_split_point_for_tool_pairs(
+                agent.messages, messages_to_summarize_count
+            )
+
+            if messages_to_summarize_count <= 0:
+                raise ContextWindowOverflowException("Cannot summarize: insufficient messages for summarization")
+
+            # Extract messages to summarize
+            messages_to_summarize = agent.messages[:messages_to_summarize_count]
+            remaining_messages = agent.messages[messages_to_summarize_count:]
+
+            # Keep track of the number of messages that have been summarized thus far.
+            self.removed_message_count += len(messages_to_summarize)
+            # If there is a summary message, don't count it in the removed_message_count.
+            if self._summary_message:
+                self.removed_message_count -= 1
+
+            # Generate summary
+            self._summary_message = self._generate_summary(messages_to_summarize, agent)
+
+            # Replace the summarized messages with the summary
+            agent.messages[:] = [self._summary_message] + remaining_messages
+
+        except Exception as summarization_error:
+            logger.error("Summarization failed: %s", summarization_error)
+            raise summarization_error from e
+
+    def _generate_summary(self, messages: List[Message], agent: "Agent") -> Message:
+        """Generate a summary of the provided messages.
+
+        Args:
+            messages: The messages to summarize.
+            agent: The agent instance to use for summarization.
+
+        Returns:
+            A message containing the conversation summary.
+
+        Raises:
+            Exception: If summary generation fails.
+        """
+        # Choose which agent to use for summarization
+        summarization_agent = self.summarization_agent if self.summarization_agent is not None else agent
+
+        # Save original system prompt and messages to restore later
+        original_system_prompt = summarization_agent.system_prompt
+        original_messages = summarization_agent.messages.copy()
+
+        try:
+            # Only override system prompt if no agent was provided during initialization
+            if self.summarization_agent is None:
+                # Use custom system prompt if provided, otherwise use default
+                system_prompt = (
+                    self.summarization_system_prompt
+                    if self.summarization_system_prompt is not None
+                    else DEFAULT_SUMMARIZATION_PROMPT
+                )
+                # Temporarily set the system prompt for summarization
+                summarization_agent.system_prompt = system_prompt
+            summarization_agent.messages = messages
+
+            # Use the agent to generate summary with rich content (can use tools if needed)
+            result = summarization_agent("Please summarize this conversation.")
+
+            return result.message
+
+        finally:
+            # Restore original agent state
+            summarization_agent.system_prompt = original_system_prompt
+            summarization_agent.messages = original_messages
+
+    def _adjust_split_point_for_tool_pairs(self, messages: List[Message], split_point: int) -> int:
+        """Adjust the split point to avoid breaking ToolUse/ToolResult pairs.
+
+        Uses the same logic as SlidingWindowConversationManager for consistency.
+
+        Args:
+            messages: The full list of messages.
+            split_point: The initially calculated split point.
+
+        Returns:
+            The adjusted split point that doesn't break ToolUse/ToolResult pairs.
+
+        Raises:
+            ContextWindowOverflowException: If no valid split point can be found.
+        """
+        if split_point > len(messages):
+            raise ContextWindowOverflowException("Split point exceeds message array length")
+
+        if split_point == len(messages):
+            return split_point
+
+        # Find the next valid split_point
+        while split_point < len(messages):
+            if (
+                # Oldest message cannot be a toolResult because it needs a toolUse preceding it
+                any("toolResult" in content for content in messages[split_point]["content"])
+                or (
+                    # Oldest message can be a toolUse only if a toolResult immediately follows it.
+                    any("toolUse" in content for content in messages[split_point]["content"])
+                    and split_point + 1 < len(messages)
+                    and not any("toolResult" in content for content in messages[split_point + 1]["content"])
+                )
+            ):
+                split_point += 1
+            else:
+                break
+        else:
+            # If we didn't find a valid split_point, then we throw
+            raise ContextWindowOverflowException("Unable to trim conversation context!")
+
+        return split_point