Skip to content

Message context manager

Message context management for agent state in TAF.

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. Optionally removes tool-related messages (AI messages with tool calls and tool result messages). 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
remove_tool_msgs
Source code in agentflow/state/message_context_manager.py 
 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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
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.
    Optionally removes tool-related messages (AI messages with tool calls and tool result messages).
    Generic over AgentState or its subclasses.
    """

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

        Args:
            max_messages (int): Maximum number of
                user messages to keep in context. Default is 10.
            remove_tool_msgs (bool): Whether to remove tool messages from context.
                Default is False.
        """
        self.max_messages = max_messages
        self.remove_tool_msgs = remove_tool_msgs
        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

        # First, remove tool messages if requested
        if self.remove_tool_msgs:
            messages = remove_tool_messages(messages)
            logger.debug("Removed tool messages, %d messages remaining", len(messages))

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

        if user_message_count <= self.max_messages:
            # Check if we removed tool messages but no trimming needed
            if self.remove_tool_msgs:
                # Return the filtered messages even if count is within limits
                return 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"]

        # Find the index of the oldest user message to keep
        user_count = 0
        start_index = len(non_system_messages)

        # Iterate from the end to find the position to start keeping messages
        for i in range(len(non_system_messages) - 1, -1, -1):
            msg = non_system_messages[i]
            if msg.role == "user":
                user_count += 1
                if user_count == self.max_messages:
                    start_index = i
                    break

        # Keep messages from start_index onwards
        final_non_system = non_system_messages[start_index:]

        # 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),
            self.max_messages,
        )

        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.

        If `remove_tool_msgs` is True, also removes:
        - AI messages that contain tool calls (intermediate tool-calling messages)
        - Tool result messages (role="tool")

        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.

        If `remove_tool_msgs` is True, also removes:
        - AI messages that contain tool calls (intermediate tool-calling messages)
        - Tool result messages (role="tool")

        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
remove_tool_msgs instance-attribute 
remove_tool_msgs = remove_tool_msgs

Functions

__init__ 
__init__(max_messages=10, remove_tool_msgs=False)

Initialize the MessageContextManager.

Parameters:

Name Type Description Default
max_messages
int

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

 10
remove_tool_msgs
bool

Whether to remove tool messages from context. Default is False.

 False
Source code in agentflow/state/message_context_manager.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
def __init__(self, max_messages: int = 10, remove_tool_msgs: bool = False) -> None:
    """
    Initialize the MessageContextManager.

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

Asynchronous version of trim_context.

If remove_tool_msgs is True, also removes: - AI messages that contain tool calls (intermediate tool-calling messages) - Tool result messages (role="tool")

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 agentflow/state/message_context_manager.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
async def atrim_context(self, state: S) -> S:
    """
    Asynchronous version of trim_context.

    If `remove_tool_msgs` is True, also removes:
    - AI messages that contain tool calls (intermediate tool-calling messages)
    - Tool result messages (role="tool")

    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.

If remove_tool_msgs is True, also removes: - AI messages that contain tool calls (intermediate tool-calling messages) - Tool result messages (role="tool")

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 agentflow/state/message_context_manager.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
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.

    If `remove_tool_msgs` is True, also removes:
    - AI messages that contain tool calls (intermediate tool-calling messages)
    - Tool result messages (role="tool")

    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

Functions