test_summarization_middleware_token_retention_preserves_ai_tool_pairs() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_summarization.py lines 346–387

def test_summarization_middleware_token_retention_preserves_ai_tool_pairs() -> None:
    """Ensure token retention preserves AI/Tool message pairs together."""

    def token_counter(messages: Iterable[MessageLikeRepresentation]) -> int:
        return sum(len(getattr(message, "content", "")) for message in messages)

    middleware = SummarizationMiddleware(
        model=ProfileChatModel(),
        trigger=("fraction", 0.1),
        keep=("fraction", 0.5),
        token_counter=token_counter,
    )

    # Total tokens: 300 + 200 + 50 + 180 + 160 = 890
    # Target keep: 500 tokens (50% of 1000)
    # Binary search finds cutoff around index 2 (ToolMessage)
    # We move back to index 1 to preserve the AIMessage with its ToolMessage
    messages: list[AnyMessage] = [
        HumanMessage(content="H" * 300),
        AIMessage(
            content="A" * 200,
            tool_calls=[{"name": "test", "args": {}, "id": "call-1"}],
        ),
        ToolMessage(content="T" * 50, tool_call_id="call-1"),
        HumanMessage(content="H" * 180),
        HumanMessage(content="H" * 160),
    ]

    state = AgentState[Any](messages=messages)
    result = middleware.before_model(state, Runtime())
    assert result is not None

    preserved_messages = result["messages"][2:]
    # We move the cutoff back to include the AIMessage with its ToolMessage
    # So we preserve messages from index 1 onward (AI + Tool + Human + Human)
    assert preserved_messages == messages[1:]

    # Verify the AI/Tool pair is preserved together
    assert isinstance(preserved_messages[0], AIMessage)
    assert preserved_messages[0].tool_calls
    assert isinstance(preserved_messages[1], ToolMessage)
    assert preserved_messages[1].tool_call_id == preserved_messages[0].tool_calls[0]["id"]