langgraph-agent-toolkit 0.9.1__tar.gz → 0.9.2__tar.gz

{langgraph_agent_toolkit-0.9.1 → langgraph_agent_toolkit-0.9.2}/CHANGELOG.md

@@ -6,6 +6,33 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to
 [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.9.2]
+
+### Fixed
+
+- `TrimMessagesMiddleware` / `TokenTrimMiddleware` no longer silently destroy
+  the model's view when the **latest turn does not fit the budget**.
+  `trim_messages` (`strategy="last"`, `start_on="human"`) returns `[]` whenever
+  the current turn exceeds the budget — e.g. a single user message larger than
+  the token budget, or a turn with more tool calls than the message budget —
+  which invoked the model with the system prompt only and made it answer with
+  the user's question dropped. Both middlewares now floor to the latest turn
+  (kept verbatim, over budget but answerable) rather than emptying it.
+  Regression tests cover an oversize user message, an oversize tool result, and
+  a 6+ tool-call burst.
+
+### Added
+
+- `TokenTrimMiddleware` — bounds the model's message view by a **token** budget
+  (the token-counting companion to `TrimMessagesMiddleware`, which bounds by
+  message count). View-only, so the full history stays in state and the system
+  prompt is preserved. The token counter is configurable (`token_counter`,
+  default `count_tokens_approximately` — no extra dependency; pass a
+  `tiktoken`-backed counter or the model for an exact count), and the budget
+  defaults to the new `DEFAULT_MAX_TOKENS_HISTORY_LENGTH` setting (unset by
+  default, since a sensible budget is model-specific). Compose it after
+  `TrimMessagesMiddleware` to cap both message count and token size.
+
 ## [0.9.1]
 
 ### Added

{langgraph_agent_toolkit-0.9.1 → langgraph_agent_toolkit-0.9.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langgraph-agent-toolkit
-Version: 0.9.1
+Version: 0.9.2
 Summary: Full toolkit for running an AI agent service built with LangGraph, FastAPI and Streamlit
 Project-URL: documentation, https://kryvokhyzha.github.io/langgraph-agent-toolkit
 Project-URL: repository, https://github.com/kryvokhyzha/langgraph-agent-toolkit

{langgraph_agent_toolkit-0.9.1 → langgraph_agent_toolkit-0.9.2}/langgraph_agent_toolkit/agents/components/middlewares/__init__.py

@@ -5,6 +5,7 @@ from langgraph_agent_toolkit.agents.components.middlewares.clear_intermediate_to
 )
 from langgraph_agent_toolkit.agents.components.middlewares.immediate_generation import ImmediateGenerationMiddleware
 from langgraph_agent_toolkit.agents.components.middlewares.sanitize_history import SanitizeHistoryMiddleware
+from langgraph_agent_toolkit.agents.components.middlewares.token_trim import TokenTrimMiddleware
 from langgraph_agent_toolkit.agents.components.middlewares.trim_messages import TrimMessagesMiddleware
 
 
@@ -12,5 +13,6 @@ __all__ = [
     "ClearIntermediateToolCallsMiddleware",
     "ImmediateGenerationMiddleware",
     "SanitizeHistoryMiddleware",
+    "TokenTrimMiddleware",
     "TrimMessagesMiddleware",
 ]

langgraph_agent_toolkit-0.9.2/langgraph_agent_toolkit/agents/components/middlewares/_history.py

@@ -0,0 +1,39 @@
+"""Internal helper shared by the history-trimming middleware (``TrimMessages`` / ``TokenTrim``)."""
+
+from langchain_core.messages import BaseMessage
+
+
+def keep_latest_turn_if_emptied(
+    original: list[BaseMessage],
+    trimmed: list[BaseMessage],
+) -> list[BaseMessage]:
+    """Guard against ``trim_messages`` destroying the model's view.
+
+    ``trim_messages(strategy="last", start_on="human", allow_partial=False)`` returns an empty (or
+    system-only) list whenever the latest human turn does not fit the budget — e.g. a single user
+    message larger than the budget, or a current turn with more tool calls than the budget allows. The
+    model would then be invoked with no user content and answer from the system prompt alone.
+
+    When trimming leaves no human/tool message, fall back to the latest turn (from the last
+    ``HumanMessage`` to the end, plus any leading system messages) verbatim, bypassing the budget for
+    that turn: a too-long-but-answerable prompt beats a hallucinated answer to nothing, and an
+    over-budget turn surfaces as a visible model context error rather than silent corruption.
+
+    Args:
+        original: The messages before trimming.
+        trimmed: The result of ``trim_messages``.
+
+    Returns:
+        ``trimmed`` when it still carries user content, else the latest turn from ``original``.
+
+    """
+    if any(m.type in ("human", "tool") for m in trimmed):
+        return trimmed
+
+    last_human = next((i for i in range(len(original) - 1, -1, -1) if original[i].type == "human"), None)
+    if last_human is None:
+        # No human turn to anchor on (pathological); preserve everything rather than wipe the view.
+        return original
+
+    leading_system = [m for m in original[:last_human] if m.type == "system"]
+    return leading_system + original[last_human:]

langgraph_agent_toolkit-0.9.2/langgraph_agent_toolkit/agents/components/middlewares/token_trim.py

@@ -0,0 +1,81 @@
+"""Middleware that trims the model's input to a bounded token budget (view-only)."""
+
+from collections.abc import Awaitable, Callable
+
+from langchain.agents.middleware import AgentMiddleware
+from langchain.agents.middleware.types import ModelRequest, ModelResponse
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.messages import BaseMessage
+from langchain_core.messages.utils import count_tokens_approximately, trim_messages
+
+from langgraph_agent_toolkit.agents.components.middlewares._history import keep_latest_turn_if_emptied
+from langgraph_agent_toolkit.core.settings import settings
+
+
+class TokenTrimMiddleware(AgentMiddleware):
+    """Bound the model's input to a token budget, non-destructively.
+
+    The token-counting companion to ``TrimMessagesMiddleware`` (which bounds by *message* count). Each
+    model call sees at most ``max_tokens`` tokens of recent history (trimmed to start on a human turn
+    and to keep the system prompt), while the full history stays in state. Compose the two to cap both
+    message count and token size, as the custom ``create_react_agent``'s ``pre_model_hook`` did.
+
+    Args:
+        max_tokens: Token budget for the model's message view. Defaults to
+            ``settings.DEFAULT_MAX_TOKENS_HISTORY_LENGTH``. ``None`` (the toolkit default) disables
+            trimming — set it, or pass ``max_tokens``, to a budget appropriate for your model.
+        token_counter: How to count tokens — a per-message or per-list callable, or a chat model.
+            Defaults to ``count_tokens_approximately`` (fast, model-agnostic, no extra dependency).
+            Pass a ``tiktoken``-backed counter (or the model itself) for an exact count.
+
+    Note:
+        On ``create_agent`` the system prompt is never part of ``request.messages`` (it lives on
+        ``request.system_message`` and is prepended at call time), so it is always preserved
+        regardless of the budget.
+
+    """
+
+    def __init__(
+        self,
+        max_tokens: int | None = None,
+        token_counter: (
+            Callable[[list[BaseMessage]], int] | Callable[[BaseMessage], int] | BaseLanguageModel
+        ) = count_tokens_approximately,
+    ) -> None:
+        super().__init__()
+        resolved = max_tokens if max_tokens is not None else settings.DEFAULT_MAX_TOKENS_HISTORY_LENGTH
+        if resolved is not None and resolved < 1:
+            raise ValueError("max_tokens must be >= 1 (or None to disable)")
+        self.max_tokens = resolved
+        self.token_counter = token_counter
+
+    def _trim(self, messages: list[BaseMessage]) -> list[BaseMessage]:
+        if not self.max_tokens:
+            return messages
+        trimmed = trim_messages(
+            messages,
+            token_counter=self.token_counter,
+            max_tokens=self.max_tokens,
+            strategy="last",
+            start_on="human",
+            end_on=("human", "tool"),
+            include_system=True,
+            allow_partial=False,
+        )
+        # A latest turn larger than the budget collapses to []; keep that turn (over budget but
+        # answerable) instead of invoking the model with no user content.
+        return keep_latest_turn_if_emptied(messages, trimmed)
+
+    def wrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], ModelResponse],
+    ) -> ModelResponse:
+        return handler(request.override(messages=self._trim(request.messages)))
+
+    async def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelResponse:
+        return await handler(request.override(messages=self._trim(request.messages)))

{langgraph_agent_toolkit-0.9.1 → langgraph_agent_toolkit-0.9.2}/langgraph_agent_toolkit/agents/components/middlewares/trim_messages.py

@@ -7,6 +7,7 @@ from langchain.agents.middleware.types import ModelRequest, ModelResponse
 from langchain_core.messages import BaseMessage
 from langchain_core.messages.utils import trim_messages
 
+from langgraph_agent_toolkit.agents.components.middlewares._history import keep_latest_turn_if_emptied
 from langgraph_agent_toolkit.core.settings import settings
 
 
@@ -32,7 +33,7 @@ class TrimMessagesMiddleware(AgentMiddleware):
         self.max_messages = resolved
 
     def _trim(self, messages: list[BaseMessage]) -> list[BaseMessage]:
-        return trim_messages(
+        trimmed = trim_messages(
             messages,
             token_counter=len,  # count messages, not tokens
             max_tokens=self.max_messages,
@@ -42,6 +43,9 @@ class TrimMessagesMiddleware(AgentMiddleware):
             include_system=True,
             allow_partial=False,
         )
+        # A current turn longer than max_messages collapses to []; keep that turn instead of
+        # invoking the model with no user content.
+        return keep_latest_turn_if_emptied(messages, trimmed)
 
     def wrap_model_call(
         self,

{langgraph_agent_toolkit-0.9.1 → langgraph_agent_toolkit-0.9.2}/langgraph_agent_toolkit/core/_base_settings.py

@@ -162,6 +162,9 @@ class Settings(BaseSettings):
     # Agent configuration
     DEFAULT_AGENT: str = "create-agent"
     DEFAULT_MAX_MESSAGE_HISTORY_LENGTH: int = 18
+    # Token budget for the model's message view (TokenTrimMiddleware). None disables it (opt-in,
+    # since a sensible budget is model-specific); set a value to bound history by tokens.
+    DEFAULT_MAX_TOKENS_HISTORY_LENGTH: int | None = None
     DEFAULT_RECURSION_LIMIT: int = 64
     MULTIMODAL_MAX_ATTACHMENTS: int | None = Field(
         default=None,

{langgraph_agent_toolkit-0.9.1 → langgraph_agent_toolkit-0.9.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph-agent-toolkit"
-version = "0.9.1"
+version = "0.9.2"
 description = "Full toolkit for running an AI agent service built with LangGraph, FastAPI and Streamlit"
 readme = "README.md"
 authors = [

{langgraph_agent_toolkit-0.9.1 → langgraph_agent_toolkit-0.9.2}/tests/agents/test_middlewares.py

@@ -11,6 +11,7 @@ from langgraph_agent_toolkit.agents.components.middlewares import (
     ClearIntermediateToolCallsMiddleware,
     ImmediateGenerationMiddleware,
     SanitizeHistoryMiddleware,
+    TokenTrimMiddleware,
     TrimMessagesMiddleware,
 )
 from langgraph_agent_toolkit.core.models.fake import FakeToolModel
@@ -246,6 +247,143 @@ def test_trim_messages_rejects_bad_max():
         TrimMessagesMiddleware(max_messages=0)
 
 
+def _content_len(messages: list) -> int:
+    """Per-list token-counter stub: total characters of message content."""
+    return sum(len(m.content) for m in messages)
+
+
+def test_token_trim_bounds_to_token_budget():
+    """The model's view is trimmed until the token count fits the budget, starting on a human turn."""
+    mw = TokenTrimMiddleware(max_tokens=12, token_counter=_content_len)
+    request = MagicMock()
+    request.messages = [
+        HumanMessage("aaaa"),
+        AIMessage("bbbb"),
+        HumanMessage("cccc"),
+        AIMessage("dddd"),
+        HumanMessage("eeee"),
+    ]
+    request.override.return_value = "TRIMMED"
+    handler = MagicMock(return_value="R")
+
+    mw.wrap_model_call(request, handler)
+
+    trimmed = request.override.call_args.kwargs["messages"]
+    assert _content_len(trimmed) <= 12
+    assert trimmed[0].type == "human"
+    handler.assert_called_once_with("TRIMMED")
+
+
+def test_token_trim_keeps_system_message():
+    """include_system keeps a system message even when the rest is trimmed to fit the budget."""
+    mw = TokenTrimMiddleware(max_tokens=8, token_counter=_content_len)
+    request = MagicMock()
+    request.messages = [SystemMessage("sys"), HumanMessage("aaaa"), AIMessage("bbbb"), HumanMessage("cccc")]
+    handler = MagicMock()
+
+    mw.wrap_model_call(request, handler)
+
+    trimmed = request.override.call_args.kwargs["messages"]
+    assert any(m.type == "system" for m in trimmed)
+
+
+def test_token_trim_disabled_passes_through():
+    """With no budget (None), the middleware is a no-op."""
+    mw = TokenTrimMiddleware(max_tokens=None)
+    assert mw.max_tokens is None  # toolkit default is opt-in
+    request = MagicMock()
+    request.messages = [HumanMessage("q1"), AIMessage("a1"), HumanMessage("q2")]
+    handler = MagicMock()
+
+    mw.wrap_model_call(request, handler)
+
+    assert request.override.call_args.kwargs["messages"] == request.messages
+
+
+def test_token_trim_default_from_settings():
+    assert TokenTrimMiddleware().max_tokens == settings.DEFAULT_MAX_TOKENS_HISTORY_LENGTH
+
+
+def test_token_trim_rejects_bad_max():
+    with pytest.raises(ValueError, match="max_tokens"):
+        TokenTrimMiddleware(max_tokens=0)
+
+
+def _latest_human_present(view: list, text: str) -> bool:
+    return any(m.type == "human" and text in str(m.content) for m in view)
+
+
+def test_token_trim_preserves_oversize_latest_human():
+    """A single user message larger than the budget is kept, not dropped to an empty/system-only view."""
+    mw = TokenTrimMiddleware(max_tokens=5, token_counter=_content_len)
+    request = MagicMock()
+    request.messages = [HumanMessage("this question is far longer than the tiny budget UNIQUE")]
+    handler = MagicMock()
+
+    mw.wrap_model_call(request, handler)
+
+    kept = request.override.call_args.kwargs["messages"]
+    assert _latest_human_present(kept, "UNIQUE")  # question preserved despite exceeding the budget
+
+
+def test_token_trim_preserves_oversize_tool_turn():
+    """A current turn (human + big tool result) over budget is kept verbatim, not emptied."""
+    mw = TokenTrimMiddleware(max_tokens=5, token_counter=_content_len)
+    request = MagicMock()
+    request.messages = [
+        HumanMessage("q UNIQUE"),
+        _ai_tool("search", "s1"),
+        ToolMessage("x" * 200, tool_call_id="s1", name="search"),
+    ]
+    handler = MagicMock()
+
+    mw.wrap_model_call(request, handler)
+
+    kept = request.override.call_args.kwargs["messages"]
+    assert _latest_human_present(kept, "UNIQUE")
+    ai_ids = {c["id"] for m in kept if m.type == "ai" and m.tool_calls for c in m.tool_calls}
+    assert all(m.tool_call_id in ai_ids for m in kept if m.type == "tool")  # pairing intact
+
+
+def test_trim_messages_preserves_long_tool_burst_turn():
+    """A current turn with more messages than max_messages is kept, not wiped to an empty view."""
+    mw = TrimMessagesMiddleware(max_messages=12)
+    burst = [HumanMessage("BURST_Q")]
+    for i in range(6):  # 1 human + 6*(AI tool_call, Tool) = 13 messages > 12
+        burst += [_ai_tool("search", f"s{i}"), ToolMessage("r", tool_call_id=f"s{i}", name="search")]
+    request = MagicMock()
+    request.messages = burst
+    handler = MagicMock()
+
+    mw.wrap_model_call(request, handler)
+
+    kept = request.override.call_args.kwargs["messages"]
+    assert _latest_human_present(kept, "BURST_Q")  # the human survives the burst
+    ai_ids = {c["id"] for m in kept if m.type == "ai" and m.tool_calls for c in m.tool_calls}
+    assert all(m.tool_call_id in ai_ids for m in kept if m.type == "tool")
+
+
+def test_keep_latest_turn_floor_keeps_only_latest_turn_and_leading_system():
+    """The floor preserves leading system messages + the latest turn, dropping older over-budget turns."""
+    from langgraph_agent_toolkit.agents.components.middlewares._history import keep_latest_turn_if_emptied
+
+    original = [SystemMessage("sys"), HumanMessage("old"), AIMessage("a"), HumanMessage("latest")]
+    floored = keep_latest_turn_if_emptied(original, [])  # [] simulates a budget collapse
+
+    assert floored[0].type == "system"
+    assert floored[-1].content == "latest"
+    assert [m.content for m in floored if m.type == "human"] == ["latest"]  # only the latest turn's human
+
+
+def test_keep_latest_turn_floor_passes_through_healthy_trim():
+    """When trimming kept user content, the floor returns it unchanged."""
+    from langgraph_agent_toolkit.agents.components.middlewares._history import keep_latest_turn_if_emptied
+
+    original = [HumanMessage("a"), AIMessage("b"), HumanMessage("c")]
+    trimmed = [HumanMessage("c")]
+    assert keep_latest_turn_if_emptied(original, trimmed) is trimmed
+
+
 def test_create_agent_with_middleware_runs_end_to_end():
     """create_agent composes both middleware and runs without error (FAKE model, no tool loop)."""