-
Notifications
You must be signed in to change notification settings - Fork 569
feat: add per_turn parameter to SlidingWindowConversationManager #1374
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -6,6 +6,7 @@ | |
| if TYPE_CHECKING: | ||
| from ...agent.agent import Agent | ||
|
|
||
| from ...hooks import BeforeModelCallEvent, HookRegistry | ||
| from ...types.content import Messages | ||
| from ...types.exceptions import ContextWindowOverflowException | ||
| from .conversation_manager import ConversationManager | ||
|
|
@@ -18,19 +19,107 @@ class SlidingWindowConversationManager(ConversationManager): | |
|
|
||
| This class handles the logic of maintaining a conversation window that preserves tool usage pairs and avoids | ||
| invalid window states. | ||
|
|
||
| Supports proactive management during agent loop execution via the per_turn parameter. | ||
| """ | ||
|
|
||
| def __init__(self, window_size: int = 40, should_truncate_results: bool = True): | ||
| def __init__(self, window_size: int = 40, should_truncate_results: bool = True, *, per_turn: bool | int = False): | ||
| """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 | ||
| per_turn: Controls when to apply message management during agent execution. | ||
| - False (default): Only apply management at the end (default behavior) | ||
| - True: Apply management before every model call | ||
| - int (e.g., 3): Apply management before every N model calls | ||
|
|
||
| When to use per_turn: If your agent performs many tool operations in loops | ||
| (e.g., web browsing with frequent screenshots), enable per_turn to proactively | ||
| manage message history and prevent the agent loop from slowing down. Start with | ||
| per_turn=True and adjust to a specific frequency (e.g., per_turn=5) if needed | ||
| for performance tuning. | ||
|
|
||
| Raises: | ||
| ValueError: If per_turn is 0 or a negative integer. | ||
| """ | ||
| super().__init__() | ||
|
|
||
| # Validate per_turn parameter | ||
| # Note: Must check bool before int since bool is a subclass of int in Python | ||
| if not isinstance(per_turn, bool) and isinstance(per_turn, int) and per_turn <= 0: | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In our SDK I don't think we validate input parameters
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Are we against it or just that this is contrary to other places? Personally I don't see a problem with enforcing it for new things, but can remove it if we want
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. While validation wouldn't hurt, we intentionally avoid it. @pgrayy mentioned it will make SDK cumbersome, and yes, if we do it here, we can do else where. |
||
| raise ValueError(f"per_turn must be True, False, or a positive integer, got {per_turn}") | ||
|
|
||
| self.window_size = window_size | ||
| self.should_truncate_results = should_truncate_results | ||
| self.per_turn = per_turn | ||
| self.model_call_count = 0 | ||
|
|
||
| def register_hooks(self, registry: "HookRegistry", **kwargs: Any) -> None: | ||
| """Register hook callbacks for per-turn conversation management. | ||
|
|
||
| Args: | ||
| registry: The hook registry to register callbacks with. | ||
| **kwargs: Additional keyword arguments for future extensibility. | ||
| """ | ||
| super().register_hooks(registry, **kwargs) | ||
|
|
||
| # Always register the callback - per_turn check happens in the callback | ||
| registry.add_callback(BeforeModelCallEvent, self._on_before_model_call) | ||
|
|
||
| def _on_before_model_call(self, event: BeforeModelCallEvent) -> None: | ||
| """Handle before model call event for per-turn management. | ||
|
|
||
| This callback is invoked before each model call. It tracks the model call count and applies message management | ||
| based on the per_turn configuration. | ||
|
|
||
| Args: | ||
| event: The before model call event containing the agent and model execution details. | ||
| """ | ||
| # Check if per_turn is enabled | ||
| if self.per_turn is False: | ||
| return | ||
|
|
||
| self.model_call_count += 1 | ||
|
|
||
| # Determine if we should apply management | ||
| should_apply = False | ||
| if self.per_turn is True: | ||
| should_apply = True | ||
| elif isinstance(self.per_turn, int) and self.per_turn > 0: | ||
| should_apply = self.model_call_count % self.per_turn == 0 | ||
|
|
||
| if should_apply: | ||
| logger.debug( | ||
| "model_call_count=<%d>, per_turn=<%s> | applying per-turn conversation management", | ||
| self.model_call_count, | ||
| self.per_turn, | ||
| ) | ||
| self.apply_management(event.agent) | ||
|
|
||
| def get_state(self) -> dict[str, Any]: | ||
| """Get the current state of the conversation manager. | ||
|
|
||
| Returns: | ||
| Dictionary containing the manager's state, including model call count for per-turn tracking. | ||
| """ | ||
| state = super().get_state() | ||
| state["model_call_count"] = self.model_call_count | ||
| return state | ||
|
|
||
| def restore_from_session(self, state: dict[str, Any]) -> Optional[list]: | ||
| """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 agent's messages. | ||
| """ | ||
| result = super().restore_from_session(state) | ||
| self.model_call_count = state.get("model_call_count", 0) | ||
| return result | ||
|
|
||
| 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. | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
session manager uses conversation manager, right? Should we put conversation_manager hooks before it in case there is something the conversation manager needs to do earlier?
Also what are you thinking for documentation? It is a bit unclear now that conversation manager is now both called manually and via hooks
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Also, what is the story for transitioning from "manual" conversation management, meaning conversation_manager.apply_management, to using only hooks.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I actually think we want to do it after, in case the conversation_manager wants to trim after the session-manager
I will have to do a docs PR for this; specifically I need to update it to account for the per_turn parameter AND conversation managers being able to register for hooks.
I'll cut an issue - though IMHO it's low priority