From 03373459dcc5f1a05dc48caeb18e2e5dc67a5cb8 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Fri, 14 Nov 2025 14:57:23 -0500 Subject: [PATCH] context trimming --- .../langchain/agents/middleware/__init__.py | 12 + .../langchain/agents/middleware/_context.py | 30 ++ .../agents/middleware/context_editing.py | 294 +++++++++++++++++- .../agents/middleware/summarization.py | 21 +- .../implementations/test_context_editing.py | 151 +++++++++ 5 files changed, 481 insertions(+), 27 deletions(-) create mode 100644 libs/langchain_v1/langchain/agents/middleware/_context.py diff --git a/libs/langchain_v1/langchain/agents/middleware/__init__.py b/libs/langchain_v1/langchain/agents/middleware/__init__.py index 44389008aa4..30d59247700 100644 --- a/libs/langchain_v1/langchain/agents/middleware/__init__.py +++ b/libs/langchain_v1/langchain/agents/middleware/__init__.py @@ -1,5 +1,12 @@ """Entrypoint to using [middleware](https://docs.langchain.com/oss/python/langchain/middleware) plugins with [Agents](https://docs.langchain.com/oss/python/langchain/agents).""" # noqa: E501 +from ._context import ( + ContextCondition, + ContextFraction, + ContextMessages, + ContextSize, + ContextTokens, +) from .context_editing import ( ClearToolUsesEdit, ContextEditingMiddleware, @@ -45,7 +52,12 @@ __all__ = [ "AgentState", "ClearToolUsesEdit", "CodexSandboxExecutionPolicy", + "ContextCondition", "ContextEditingMiddleware", + "ContextFraction", + "ContextMessages", + "ContextSize", + "ContextTokens", "DockerExecutionPolicy", "FilesystemFileSearchMiddleware", "HostExecutionPolicy", diff --git a/libs/langchain_v1/langchain/agents/middleware/_context.py b/libs/langchain_v1/langchain/agents/middleware/_context.py new file mode 100644 index 00000000000..59a0dc95bc8 --- /dev/null +++ b/libs/langchain_v1/langchain/agents/middleware/_context.py @@ -0,0 +1,30 @@ +"""Shared context size types for middleware that manages message history.""" + +from typing import Literal + +ContextFraction = tuple[Literal["fraction"], float] +"""Tuple specifying context size as a fraction of the model's context window.""" + +ContextTokens = tuple[Literal["tokens"], int] +"""Tuple specifying context size as a number of tokens.""" + +ContextMessages = tuple[Literal["messages"], int] +"""Tuple specifying context size as a number of messages.""" + +ContextSize = ContextFraction | ContextTokens | ContextMessages +"""Context size tuple to specify how much history to preserve or trigger conditions.""" + +ContextCondition = ContextSize | list[ContextSize | list[ContextSize]] +"""Recursive type to support nested AND/OR conditions. + +Top-level list = OR logic, nested list = AND logic. +""" + + +__all__ = [ + "ContextCondition", + "ContextFraction", + "ContextMessages", + "ContextSize", + "ContextTokens", +] diff --git a/libs/langchain_v1/langchain/agents/middleware/context_editing.py b/libs/langchain_v1/langchain/agents/middleware/context_editing.py index 29fcd991546..5ed123c0d3e 100644 --- a/libs/langchain_v1/langchain/agents/middleware/context_editing.py +++ b/libs/langchain_v1/langchain/agents/middleware/context_editing.py @@ -9,9 +9,9 @@ chat model. from __future__ import annotations -from collections.abc import Awaitable, Callable, Iterable, Sequence +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass -from typing import Literal +from typing import TYPE_CHECKING, Any, Literal from langchain_core.messages import ( AIMessage, @@ -30,6 +30,9 @@ from langchain.agents.middleware.types import ( ModelResponse, ) +if TYPE_CHECKING: + from langchain.agents.middleware._context import ContextCondition, ContextSize + DEFAULT_TOOL_PLACEHOLDER = "[cleared]" @@ -54,16 +57,37 @@ class ContextEdit(Protocol): @dataclass(slots=True) class ClearToolUsesEdit(ContextEdit): - """Configuration for clearing tool outputs when token limits are exceeded.""" + """Configuration for clearing tool outputs when token limits are exceeded. - trigger: int = 100_000 - """Token count that triggers the edit.""" + Supports flexible trigger and keep configurations using `ContextSize` tuples or + backwards-compatible integer values. + """ + + trigger: ContextCondition | int = 100_000 + """Trigger condition(s) for when the edit should run. + + Supports flexible AND/OR logic via nested lists: + - Single condition: `("messages", 50)` or `("tokens", 3000)` + - OR conditions: `[("tokens", 3000), ("messages", 100)]` + - AND conditions: `[("tokens", 500), ("fraction", 0.8)]` as nested list + - Mixed AND/OR: `[("messages", 10), [("tokens", 500), ("fraction", 0.8)]]` + + For backwards compatibility, also accepts an integer token count. + """ clear_at_least: int = 0 """Minimum number of tokens to reclaim when the edit runs.""" - keep: int = 3 - """Number of most recent tool results that must be preserved.""" + keep: ContextSize | int = 3 + """Context retention policy for tool results. + + Provide a `ContextSize` tuple to specify how much history to preserve: + - `("messages", 3)` - Keep last 3 tool results + - `("tokens", 1000)` - Keep tool results within token budget + - `("fraction", 0.3)` - Keep tool results within 30% of model's max tokens + + For backwards compatibility, also accepts an integer message count. + """ clear_tool_inputs: bool = False """Whether to clear the originating tool call parameters on the AI message.""" @@ -74,6 +98,66 @@ class ClearToolUsesEdit(ContextEdit): placeholder: str = DEFAULT_TOOL_PLACEHOLDER """Placeholder text inserted for cleared tool outputs.""" + model: Any = None + """Optional model instance for fractional token limits.""" + + _trigger_conditions: list[ContextSize | list[ContextSize]] | None = None + _keep_normalized: ContextSize | None = None + _trigger_as_int: int | None = None + _keep_as_int: int | None = None + + def __post_init__(self) -> None: + """Validate and normalize trigger/keep parameters.""" + # Normalize trigger + if isinstance(self.trigger, int): + self._trigger_as_int = self.trigger + self._trigger_conditions = None + elif isinstance(self.trigger, tuple): + # Single ContextSize + self._validate_context_size(self.trigger, "trigger") + self._trigger_conditions = [self.trigger] + self._trigger_as_int = None + elif isinstance(self.trigger, list): + # List of conditions + self._trigger_conditions = self._validate_trigger_conditions(self.trigger) + self._trigger_as_int = None + else: + msg = f"trigger must be int or ContextCondition, got {type(self.trigger).__name__}" + raise TypeError(msg) + + # Normalize keep + if isinstance(self.keep, int): + self._keep_as_int = self.keep + self._keep_normalized = None + elif isinstance(self.keep, tuple): + self._validate_context_size(self.keep, "keep") + self._keep_normalized = self.keep + self._keep_as_int = None + else: + msg = f"keep must be int or ContextSize, got {type(self.keep).__name__}" + raise TypeError(msg) + + # Check if model profile is required + requires_profile = False + if self._trigger_conditions: + requires_profile = self._requires_profile(self._trigger_conditions) + if self._keep_normalized and self._keep_normalized[0] == "fraction": + requires_profile = True + + if requires_profile and self.model is None: + msg = ( + "model parameter is required when using fractional token limits. " + "Pass a model instance or use absolute token/message counts instead." + ) + raise ValueError(msg) + + if requires_profile and self._get_profile_limits() is None: + msg = ( + "Model profile information is required to use fractional token limits. " + 'pip install "langchain[model-profiles]" or use absolute token counts instead.' + ) + raise ValueError(msg) + def apply( self, messages: list[AnyMessage], @@ -83,17 +167,20 @@ class ClearToolUsesEdit(ContextEdit): """Apply the clear-tool-uses strategy.""" tokens = count_tokens(messages) - if tokens <= self.trigger: + if not self._should_trigger(messages, tokens): return candidates = [ (idx, msg) for idx, msg in enumerate(messages) if isinstance(msg, ToolMessage) ] - if self.keep >= len(candidates): + # Determine how many tool results to keep + keep_count = self._determine_keep_count(messages, tokens) + + if keep_count >= len(candidates): candidates = [] - elif self.keep: - candidates = candidates[: -self.keep] + elif keep_count > 0: + candidates = candidates[:-keep_count] cleared_tokens = 0 excluded_tools = set(self.exclude_tools) @@ -181,6 +268,191 @@ class ClearToolUsesEdit(ContextEdit): } ) + def _should_trigger(self, messages: list[AnyMessage], total_tokens: int) -> bool: + """Determine whether the edit should trigger based on current state.""" + # Backwards compatibility: int trigger + if self._trigger_as_int is not None: + return total_tokens > self._trigger_as_int + + # New API: ContextCondition with AND/OR logic + if not self._trigger_conditions: + return False + + # OR logic across top-level conditions + for condition in self._trigger_conditions: + if isinstance(condition, list): + # AND group - all must be satisfied + if self._check_and_group(condition, messages, total_tokens): + return True + elif self._check_single_condition(condition, messages, total_tokens): + # Single condition + return True + return False + + def _check_and_group( + self, and_group: list[ContextSize], messages: list[AnyMessage], total_tokens: int + ) -> bool: + """Check if all conditions in an AND group are satisfied.""" + for condition in and_group: + if not self._check_single_condition(condition, messages, total_tokens): + return False + return True + + def _check_single_condition( + self, condition: ContextSize, messages: list[AnyMessage], total_tokens: int + ) -> bool: + """Check if a single condition is satisfied.""" + kind, value = condition + if kind == "messages": + return len(messages) >= value + if kind == "tokens": + return total_tokens >= value + if kind == "fraction": + max_input_tokens = self._get_profile_limits() + if max_input_tokens is None: + return False + threshold = int(max_input_tokens * value) + if threshold <= 0: + threshold = 1 + return total_tokens >= threshold + return False + + def _determine_keep_count(self, messages: list[AnyMessage], total_tokens: int) -> int: # noqa: ARG002 + """Determine how many tool results to keep based on keep configuration.""" + # Backwards compatibility: int keep + if self._keep_as_int is not None: + return self._keep_as_int + + # New API: ContextSize + if self._keep_normalized is None: + return 0 + + kind, value = self._keep_normalized + if kind == "messages": + return int(value) + if kind in {"tokens", "fraction"}: + # For token-based keep, we need to count backwards through tool messages + # to find how many fit within the budget + target_tokens = self._get_target_token_count(value, kind) + if target_tokens is None: + return 0 + return self._count_tool_messages_within_budget(messages, target_tokens) + return 0 + + def _get_target_token_count(self, value: float, kind: str) -> int | None: + """Get the target token count for token/fraction-based keep.""" + if kind == "fraction": + max_input_tokens = self._get_profile_limits() + if max_input_tokens is None: + return None + target = int(max_input_tokens * value) + elif kind == "tokens": + target = int(value) + else: + return None + + return max(1, target) if target > 0 else 1 + + def _count_tool_messages_within_budget( + self, messages: list[AnyMessage], target_tokens: int + ) -> int: + """Count how many recent tool messages fit within token budget.""" + tool_messages = [msg for msg in messages if isinstance(msg, ToolMessage)] + if not tool_messages: + return 0 + + # Count backwards from the end + count = 0 + accumulated_tokens = 0 + for tool_msg in reversed(tool_messages): + # Approximate token count for this message + msg_tokens = len(str(tool_msg.content)) + if accumulated_tokens + msg_tokens > target_tokens and count > 0: + break + accumulated_tokens += msg_tokens + count += 1 + + return count + + def _get_profile_limits(self) -> int | None: + """Retrieve max input token limit from the model profile.""" + if self.model is None: + return None + + try: + profile = self.model.profile + except (AttributeError, ImportError): + return None + + if not isinstance(profile, Mapping): + return None + + max_input_tokens = profile.get("max_input_tokens") + + if not isinstance(max_input_tokens, int): + return None + + return max_input_tokens + + def _validate_context_size(self, context: ContextSize, parameter_name: str) -> ContextSize: + """Validate context configuration tuples.""" + kind, value = context + if kind == "fraction": + if not 0 < value <= 1: + msg = f"Fractional {parameter_name} values must be between 0 and 1, got {value}." + raise ValueError(msg) + elif kind in {"tokens", "messages"}: + # For trigger, value must be > 0. For keep, value can be >= 0 (0 means keep nothing) + if parameter_name == "trigger" and value <= 0: + msg = f"{parameter_name} thresholds must be greater than 0, got {value}." + raise ValueError(msg) + if parameter_name == "keep" and value < 0: + msg = f"{parameter_name} values must be non-negative, got {value}." + raise ValueError(msg) + else: + msg = f"Unsupported context size type {kind} for {parameter_name}." + raise ValueError(msg) + return context + + def _validate_trigger_conditions( + self, conditions: list[Any] + ) -> list[ContextSize | list[ContextSize]]: + """Validate and normalize trigger conditions with nested AND/OR logic. + + Args: + conditions: List of ContextSize tuples or nested lists of ContextSize tuples. + + Returns: + Validated list where top-level items are OR'd and nested lists are AND'd. + """ + validated: list[ContextSize | list[ContextSize]] = [] + for item in conditions: + if isinstance(item, tuple): + # Single condition (tuple) + validated.append(self._validate_context_size(item, "trigger")) + elif isinstance(item, list): + # AND group (nested list) + if not item: + msg = "Empty AND groups are not allowed in trigger conditions." + raise ValueError(msg) + and_group = [self._validate_context_size(cond, "trigger") for cond in item] + validated.append(and_group) + else: + msg = f"Trigger conditions must be tuples or lists, got {type(item).__name__}." + raise ValueError(msg) + return validated + + def _requires_profile(self, conditions: list[ContextSize | list[ContextSize]]) -> bool: + """Check if any condition requires model profile information.""" + for condition in conditions: + if isinstance(condition, list): + # AND group + if any(c[0] == "fraction" for c in condition): + return True + elif condition[0] == "fraction": + return True + return False + class ContextEditingMiddleware(AgentMiddleware): """Automatically prune tool results to manage context size. diff --git a/libs/langchain_v1/langchain/agents/middleware/summarization.py b/libs/langchain_v1/langchain/agents/middleware/summarization.py index f6d7a2402f7..f7a9ae90150 100644 --- a/libs/langchain_v1/langchain/agents/middleware/summarization.py +++ b/libs/langchain_v1/langchain/agents/middleware/summarization.py @@ -3,7 +3,7 @@ import uuid import warnings from collections.abc import Callable, Iterable, Mapping -from typing import Any, Literal, cast +from typing import Any, cast from langchain_core.messages import ( AIMessage, @@ -19,6 +19,10 @@ from langgraph.graph.message import ( ) from langgraph.runtime import Runtime +from langchain.agents.middleware._context import ( + ContextCondition, + ContextSize, +) from langchain.agents.middleware.types import AgentMiddleware, AgentState from langchain.chat_models import BaseChatModel, init_chat_model @@ -57,21 +61,6 @@ _DEFAULT_TRIM_TOKEN_LIMIT = 4000 _DEFAULT_FALLBACK_MESSAGE_COUNT = 15 _SEARCH_RANGE_FOR_TOOL_PAIRS = 5 -ContextFraction = tuple[Literal["fraction"], float] -"""Tuple specifying context size as a fraction of the model's context window.""" -ContextTokens = tuple[Literal["tokens"], int] -"""Tuple specifying context size as a number of tokens.""" -ContextMessages = tuple[Literal["messages"], int] -"""Tuple specifying context size as a number of messages.""" - -ContextSize = ContextFraction | ContextTokens | ContextMessages -"""Context size tuple to specify how much history to preserve.""" - -ContextCondition = ContextSize | list[ContextSize | list[ContextSize]] -"""Recursive type to support nested AND/OR conditions - -Top-level list = OR logic, nested list = AND logic.""" - class SummarizationMiddleware(AgentMiddleware): """Summarizes conversation history when token limits are approached. diff --git a/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_context_editing.py b/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_context_editing.py index 7a9d901b951..38756ba9332 100644 --- a/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_context_editing.py +++ b/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_context_editing.py @@ -399,3 +399,154 @@ async def test_exclude_tools_prevents_clearing_async() -> None: assert isinstance(calc_tool, ToolMessage) assert calc_tool.content == "[cleared]" + + +# New API tests + + +def test_new_api_trigger_with_context_size_tuple() -> None: + """Test new API with ContextSize tuple for trigger.""" + tool_call_id = "call-1" + ai_message = AIMessage( + content="", + tool_calls=[{"id": tool_call_id, "name": "search", "args": {}}], + ) + tool_message = ToolMessage(content="x" * 200, tool_call_id=tool_call_id) + + state, request = _make_state_and_request([ai_message, tool_message]) + + # Use new API with tuple + edit = ClearToolUsesEdit( + trigger=("tokens", 50), + keep=("messages", 0), + placeholder="[cleared]", + ) + middleware = ContextEditingMiddleware(edits=[edit]) + + def mock_handler(req: ModelRequest) -> AIMessage: + return AIMessage(content="mock response") + + middleware.wrap_model_call(request, mock_handler) + + cleared_tool = request.messages[1] + assert isinstance(cleared_tool, ToolMessage) + assert cleared_tool.content == "[cleared]" + + +def test_new_api_keep_with_messages_tuple() -> None: + """Test new API with messages-based keep.""" + conversation: list[AIMessage | ToolMessage] = [] + for i in range(5): + call_id = f"call-{i}" + conversation.append( + AIMessage( + content="", + tool_calls=[{"id": call_id, "name": "tool", "args": {}}], + ) + ) + conversation.append(ToolMessage(content="x" * 50, tool_call_id=call_id)) + + state, request = _make_state_and_request(conversation) + + edit = ClearToolUsesEdit( + trigger=("tokens", 50), + keep=("messages", 2), # Keep last 2 tool results + placeholder="[cleared]", + ) + middleware = ContextEditingMiddleware(edits=[edit]) + + def mock_handler(req: ModelRequest) -> AIMessage: + return AIMessage(content="mock response") + + middleware.wrap_model_call(request, mock_handler) + + # Check that first 3 tool messages are cleared, last 2 are preserved + tool_messages = [msg for msg in request.messages if isinstance(msg, ToolMessage)] + cleared = [msg for msg in tool_messages if msg.content == "[cleared]"] + preserved = [msg for msg in tool_messages if msg.content != "[cleared]"] + + assert len(cleared) == 3 + assert len(preserved) == 2 + + +def test_new_api_or_conditions() -> None: + """Test new API with OR trigger conditions.""" + tool_call_id = "call-1" + ai_message = AIMessage( + content="", + tool_calls=[{"id": tool_call_id, "name": "search", "args": {}}], + ) + tool_message = ToolMessage(content="x" * 200, tool_call_id=tool_call_id) + + state, request = _make_state_and_request([ai_message, tool_message]) + + # Use OR conditions: triggers if tokens >= 50 OR messages >= 100 + edit = ClearToolUsesEdit( + trigger=[("tokens", 50), ("messages", 100)], + keep=("messages", 0), + placeholder="[cleared]", + ) + middleware = ContextEditingMiddleware(edits=[edit]) + + def mock_handler(req: ModelRequest) -> AIMessage: + return AIMessage(content="mock response") + + middleware.wrap_model_call(request, mock_handler) + + # Should trigger because tokens >= 50 (even though messages < 100) + cleared_tool = request.messages[1] + assert isinstance(cleared_tool, ToolMessage) + assert cleared_tool.content == "[cleared]" + + +def test_new_api_backwards_compatibility() -> None: + """Test that old integer API still works.""" + tool_call_id = "call-1" + ai_message = AIMessage( + content="", + tool_calls=[{"id": tool_call_id, "name": "search", "args": {}}], + ) + tool_message = ToolMessage(content="x" * 200, tool_call_id=tool_call_id) + + state, request = _make_state_and_request([ai_message, tool_message]) + + # Old API with integers + edit = ClearToolUsesEdit( + trigger=50, # int + keep=0, # int + placeholder="[cleared]", + ) + middleware = ContextEditingMiddleware(edits=[edit]) + + def mock_handler(req: ModelRequest) -> AIMessage: + return AIMessage(content="mock response") + + middleware.wrap_model_call(request, mock_handler) + + cleared_tool = request.messages[1] + assert isinstance(cleared_tool, ToolMessage) + assert cleared_tool.content == "[cleared]" + + +def test_new_api_validation_errors() -> None: + """Test that validation errors are raised for invalid configurations.""" + # Test invalid fraction value + try: + ClearToolUsesEdit(trigger=("fraction", 1.5), keep=("messages", 3)) + assert False, "Should have raised ValueError" + except ValueError as e: + assert "must be between 0 and 1" in str(e) + + # Test invalid token count + try: + ClearToolUsesEdit(trigger=("tokens", -1), keep=("messages", 3)) + assert False, "Should have raised ValueError" + except ValueError as e: + assert "must be greater than 0" in str(e) + + # Test unsupported type + try: + ClearToolUsesEdit(trigger=("invalid", 100), keep=("messages", 3)) # type: ignore[arg-type] + assert False, "Should have raised ValueError" + except ValueError as e: + assert "Unsupported context size type" in str(e)