Skip to content

Message context manager

Message context management for agent state in PyAgenity.

This module provides MessageContextManager, which trims and manages the message history (context) for agent interactions, ensuring efficient context window usage.

Classes:

Name Description
MessageContextManager

Manages the context field for AI interactions.

Attributes:

Name Type Description
S
logger

Attributes

S module-attribute 

S = TypeVar('S', bound=AgentState)

logger module-attribute 

logger = getLogger(__name__)

Classes

MessageContextManager

Bases: BaseContextManager[S]

Manages the context field for AI interactions.

This class trims the context (message history) based on a maximum number of user messages, ensuring the first message (usually a system prompt) is always preserved. Generic over AgentState or its subclasses.

Methods:

Name Description
__init__

Initialize the MessageContextManager.
atrim_context

Asynchronous version of trim_context.
trim_context

Trim the context in the given AgentState based on the maximum number of user messages.

Attributes:

Name Type Description
max_messages
Source code in pyagenity/state/message_context_manager.py 
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
class MessageContextManager(BaseContextManager[S]):
    """
    Manages the context field for AI interactions.

    This class trims the context (message history) based on a maximum number of user messages,
    ensuring the first message (usually a system prompt) is always preserved.
    Generic over AgentState or its subclasses.
    """

    def __init__(self, max_messages: int = 10) -> None:
        """
        Initialize the MessageContextManager.

        Args:
            max_messages (int): Maximum number of
                user messages to keep in context. Default is 10.
        """
        self.max_messages = max_messages
        logger.debug("Initialized MessageContextManager with max_messages=%d", max_messages)

    def _trim(self, messages: list[Message]) -> list[Message] | None:
        """
        Trim messages keeping system messages and most recent user messages.

        Returns None if no trimming is needed, otherwise returns the trimmed list.
        """
        # check context is empty
        if not messages:
            logger.debug("No messages to trim; context is empty")
            return None

        # Count user messages
        user_message_count = sum(1 for msg in messages if msg.role == "user")

        if user_message_count <= self.max_messages:
            # no trimming needed
            logger.debug(
                "No trimming needed; context is within limits (%d user messages)",
                user_message_count,
            )
            return None

        # Separate system messages (usually at the beginning)
        system_messages = [msg for msg in messages if msg.role == "system"]
        non_system_messages = [msg for msg in messages if msg.role != "system"]

        # Keep only the most recent messages that include max_messages user messages
        final_non_system = []
        user_count = 0

        # Iterate from the end to keep most recent messages
        for msg in reversed(non_system_messages):
            if msg.role == "user":
                if user_count >= self.max_messages:
                    break
                user_count += 1
            final_non_system.insert(0, msg)  # Insert at beginning to maintain order

        # Combine system messages (at start) with trimmed conversation
        trimmed_messages = system_messages + final_non_system

        logger.debug(
            "Trimmed from %d to %d messages (%d user messages kept)",
            len(messages),
            len(trimmed_messages),
            user_count,
        )

        return trimmed_messages

    def trim_context(self, state: S) -> S:
        """
        Trim the context in the given AgentState based on the maximum number of user messages.

        The first message (typically a system prompt) is always preserved. Only the most recent
        user messages up to `max_messages` are kept, along with the first message.

        Args:
            state (AgentState): The agent state containing the context to trim.

        Returns:
            S: The updated agent state with trimmed context.
        """
        messages = state.context
        trimmed_messages = self._trim(messages)
        if trimmed_messages is not None:
            state.context = trimmed_messages
        return state

    async def atrim_context(self, state: S) -> S:
        """
        Asynchronous version of trim_context.

        Args:
            state (AgentState): The agent state containing the context to trim.

        Returns:
            S: The updated agent state with trimmed context.
        """
        messages = state.context
        trimmed_messages = self._trim(messages)
        if trimmed_messages is not None:
            state.context = trimmed_messages
        return state

Attributes

max_messages instance-attribute 
max_messages = max_messages

Functions

__init__ 
__init__(max_messages=10)

Initialize the MessageContextManager.

Parameters:

Name Type Description Default
max_messages
int

Maximum number of user messages to keep in context. Default is 10.

 10
Source code in pyagenity/state/message_context_manager.py 
31
32
33
34
35
36
37
38
39
40
def __init__(self, max_messages: int = 10) -> None:
    """
    Initialize the MessageContextManager.

    Args:
        max_messages (int): Maximum number of
            user messages to keep in context. Default is 10.
    """
    self.max_messages = max_messages
    logger.debug("Initialized MessageContextManager with max_messages=%d", max_messages)
atrim_context async 
atrim_context(state)

Asynchronous version of trim_context.

Parameters:

Name Type Description Default
state
AgentState

The agent state containing the context to trim.

 required

Returns:

Name Type Description
S S

The updated agent state with trimmed context.
Source code in pyagenity/state/message_context_manager.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
async def atrim_context(self, state: S) -> S:
    """
    Asynchronous version of trim_context.

    Args:
        state (AgentState): The agent state containing the context to trim.

    Returns:
        S: The updated agent state with trimmed context.
    """
    messages = state.context
    trimmed_messages = self._trim(messages)
    if trimmed_messages is not None:
        state.context = trimmed_messages
    return state
trim_context 
trim_context(state)

Trim the context in the given AgentState based on the maximum number of user messages.

The first message (typically a system prompt) is always preserved. Only the most recent user messages up to max_messages are kept, along with the first message.

Parameters:

Name Type Description Default
state
AgentState

The agent state containing the context to trim.

 required

Returns:

Name Type Description
S S

The updated agent state with trimmed context.
Source code in pyagenity/state/message_context_manager.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def trim_context(self, state: S) -> S:
    """
    Trim the context in the given AgentState based on the maximum number of user messages.

    The first message (typically a system prompt) is always preserved. Only the most recent
    user messages up to `max_messages` are kept, along with the first message.

    Args:
        state (AgentState): The agent state containing the context to trim.

    Returns:
        S: The updated agent state with trimmed context.
    """
    messages = state.context
    trimmed_messages = self._trim(messages)
    if trimmed_messages is not None:
        state.context = trimmed_messages
    return state