mirror of
https://github.com/langchain-ai/deepagents.git
synced 2026-07-22 01:35:28 -04:00
fix(sdk): don't swallow TypeError from custom token_counter (#3927)
`SummarizationMiddleware` decides when to compact a conversation by asking a `token_counter` how large the history is. It tries to call that counter *with* the tool definitions, and previously wrapped the call in a `try/except TypeError` that retried without tools on failure. The catch couldn't tell "this counter's signature doesn't accept `tools`" apart from "this counter has a real bug" — so a `TypeError` raised inside a user-supplied counter (a tokenizer choking on an odd tool schema, a `None` where a string was expected, etc.) was silently swallowed and retried in a degraded mode that returned an undercount. The agent would then never compact and overflow the context window, or compact at the wrong moment, with no error, log, or traceback pointing at the actual culprit — turning a one-line bug into hours of wrong-tree debugging. ## Changes - Replace the call-and-catch probe in `_count_tokens` with explicit signature introspection: a new `_token_counter_accepts_tools` helper inspects the counter once and returns whether it declares a `tools` parameter (or `**kwargs`). - Resolve that decision a single time in `SummarizationMiddleware.__init__` and store it, so the per-call counting path — the hot path optimized in #3877 — pays no `inspect` cost. - Counters that accept `tools` are now called with them and any in-body `TypeError` propagates loudly; counters that clearly don't are called without them; only genuinely un-introspectable callables (e.g. C-level) fall back to the old defensive probe. The default `count_tokens_approximately` declares `tools`, so its behavior is unchanged.
This commit is contained in:
@@ -50,6 +50,7 @@ log of all evicted messages.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import logging
|
||||
import uuid
|
||||
import warnings
|
||||
@@ -184,6 +185,41 @@ class SummarizationDefaults(TypedDict):
|
||||
truncate_args_settings: TruncateArgsSettings
|
||||
|
||||
|
||||
def _token_counter_accepts_tools(counter: TokenCounter) -> bool | None:
|
||||
"""Determine whether `counter` accepts a `tools` keyword argument.
|
||||
|
||||
The `TokenCounter` contract only requires accepting messages, but the
|
||||
default counter (and most modern ones) also accept `tools=` so tool schemas
|
||||
contribute to the count. Rather than probe by calling and catching
|
||||
`TypeError` — which cannot distinguish a signature that rejects `tools`
|
||||
from a genuine `TypeError` raised inside the counter's body — the signature
|
||||
is inspected directly.
|
||||
|
||||
Args:
|
||||
counter: The token-counting callable to inspect.
|
||||
|
||||
Returns:
|
||||
`True` if the signature declares a `tools` parameter or accepts
|
||||
arbitrary keyword arguments (`**kwargs`), `False` if it clearly does
|
||||
not, or `None` when the signature cannot be introspected (some C-level
|
||||
callables expose no signature), signaling that callers should fall back
|
||||
to probing.
|
||||
"""
|
||||
try:
|
||||
parameters = inspect.signature(counter).parameters
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
for param in parameters.values():
|
||||
if param.kind is inspect.Parameter.VAR_KEYWORD:
|
||||
return True
|
||||
if param.name == "tools" and param.kind in (
|
||||
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
inspect.Parameter.KEYWORD_ONLY,
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def compute_summarization_defaults(model: BaseChatModel) -> SummarizationDefaults:
|
||||
"""Compute default summarization settings based on model profile.
|
||||
|
||||
@@ -327,6 +363,12 @@ class _DeepAgentsSummarizationMiddleware(AgentMiddleware):
|
||||
**deprecated_kwargs,
|
||||
)
|
||||
|
||||
# Whether the configured token counter accepts a `tools` kwarg. Resolved
|
||||
# once here (the counter is fixed after construction) so the per-call
|
||||
# token count never pays signature-introspection cost. `None` means the
|
||||
# signature could not be introspected, so `_count_tokens` probes instead.
|
||||
self._counter_accepts_tools = _token_counter_accepts_tools(self.token_counter)
|
||||
|
||||
# Deep Agents specific attributes
|
||||
self._backend = backend
|
||||
|
||||
@@ -725,11 +767,28 @@ A condensed summary follows:
|
||||
tools: Optional tools whose schemas contribute to the count.
|
||||
|
||||
Returns:
|
||||
Total token count. Falls back to counting without `tools` when the
|
||||
configured `token_counter` does not accept a `tools` keyword.
|
||||
Total token count. Counts without `tools` when the configured
|
||||
`token_counter` does not accept a `tools` keyword. When the
|
||||
counter's signature is introspectable, a `TypeError` raised
|
||||
inside the counter's own body is never masked — it propagates so
|
||||
a broken counter is not hidden behind a silently wrong count.
|
||||
Counters whose signature cannot be introspected are probed
|
||||
instead, and only there does a `TypeError` fall back to counting
|
||||
without `tools`.
|
||||
"""
|
||||
counted_messages = [system_message, *messages] if system_message is not None else messages
|
||||
if self._counter_accepts_tools is True:
|
||||
# `tools=` is absent from the `TokenCounter` protocol but accepted
|
||||
# here: the signature check above confirmed the counter takes it.
|
||||
return self.token_counter(counted_messages, tools=tools) # ty: ignore[unknown-argument]
|
||||
if self._counter_accepts_tools is False:
|
||||
return self.token_counter(counted_messages)
|
||||
# Signature could not be introspected; probe defensively. This is the
|
||||
# only path that swallows a `TypeError`, and only for counters whose
|
||||
# signature is opaque (some C-level callables expose no signature).
|
||||
try:
|
||||
# `tools=` is outside the `TokenCounter` protocol; the probe verifies
|
||||
# acceptance at runtime, falling back below if it is rejected.
|
||||
return self.token_counter(counted_messages, tools=tools) # ty: ignore[unknown-argument]
|
||||
except TypeError:
|
||||
return self.token_counter(counted_messages)
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
"""Unit tests for `SummarizationMiddleware` with backend offloading."""
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import time
|
||||
from collections.abc import Generator
|
||||
from collections.abc import Callable, Generator
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
from unittest.mock import MagicMock, patch
|
||||
@@ -13,7 +14,10 @@ from langchain_core.exceptions import ContextOverflowError
|
||||
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
|
||||
|
||||
from deepagents.backends.protocol import BackendProtocol, EditResult, FileDownloadResponse, ReadResult, WriteResult
|
||||
from deepagents.middleware.summarization import SummarizationMiddleware
|
||||
from deepagents.middleware.summarization import (
|
||||
SummarizationMiddleware,
|
||||
_token_counter_accepts_tools,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain.agents.middleware.types import AgentState
|
||||
@@ -2682,3 +2686,164 @@ class TestTokenCountingEfficiency:
|
||||
first_ai = captured_request.messages[0]
|
||||
assert isinstance(first_ai, AIMessage)
|
||||
assert first_ai.tool_calls[0]["args"]["content"] == "x" * 20 + "...(argument truncated)"
|
||||
|
||||
|
||||
class _OpaqueCounter:
|
||||
"""Wraps a counter so its signature cannot be introspected.
|
||||
|
||||
Accessing `__signature__` raises, which is what `inspect.signature` consults
|
||||
first. This forces `_token_counter_accepts_tools` onto its `None` result and
|
||||
`_count_tokens` onto the runtime-probe fallback -- the only path that still
|
||||
swallows a `TypeError`. The wrapped callable is invoked verbatim, so it can
|
||||
either accept or reject a `tools` kwarg at call time.
|
||||
"""
|
||||
|
||||
def __init__(self, fn: Callable[..., int]) -> None:
|
||||
self._fn = fn
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> int:
|
||||
return self._fn(*args, **kwargs)
|
||||
|
||||
@property
|
||||
def __signature__(self) -> inspect.Signature:
|
||||
msg = "no signature available"
|
||||
raise ValueError(msg)
|
||||
|
||||
|
||||
class TestTokenCounterToolsProbe:
|
||||
"""`_count_tokens` must not mask a `TypeError` raised inside the counter body."""
|
||||
|
||||
def test_accepts_tools_with_explicit_param(self) -> None:
|
||||
# The parameter must be named `tools` to exercise the signature probe.
|
||||
def counter(_messages: list[BaseMessage], *, tools: Any = None) -> int: # noqa: ANN401, ARG001
|
||||
return 0
|
||||
|
||||
assert _token_counter_accepts_tools(counter) is True
|
||||
|
||||
def test_accepts_tools_with_var_keyword(self) -> None:
|
||||
def counter(_messages: list[BaseMessage], **_kwargs: Any) -> int:
|
||||
return 0
|
||||
|
||||
assert _token_counter_accepts_tools(counter) is True
|
||||
|
||||
def test_rejects_tools_without_param(self) -> None:
|
||||
def counter(_messages: list[BaseMessage]) -> int:
|
||||
return 0
|
||||
|
||||
assert _token_counter_accepts_tools(counter) is False
|
||||
|
||||
def test_rejects_var_positional_only(self) -> None:
|
||||
# `*args` captures positional arguments only, so `tools=` cannot reach it
|
||||
# as a keyword. The counter is treated as not accepting tools, guarding
|
||||
# the deliberate omission of `VAR_POSITIONAL` from the kind check.
|
||||
def counter(_messages: list[BaseMessage], *_args: Any) -> int:
|
||||
return 0
|
||||
|
||||
assert _token_counter_accepts_tools(counter) is False
|
||||
|
||||
def test_returns_none_for_uninspectable_signature(self) -> None:
|
||||
# A callable whose signature cannot be read yields `None`, signaling that
|
||||
# `_count_tokens` must fall back to runtime probing.
|
||||
opaque = _OpaqueCounter(lambda *_args, **_kwargs: 0)
|
||||
assert _token_counter_accepts_tools(opaque) is None
|
||||
|
||||
def _make_request_with_tools(self) -> ModelRequest:
|
||||
"""A request carrying tools so `tools=` reaches the token counter."""
|
||||
state = cast("AgentState[Any]", {"messages": make_conversation_messages()})
|
||||
return ModelRequest(
|
||||
model=make_mock_model(),
|
||||
messages=state["messages"],
|
||||
system_message=None,
|
||||
tools=[{"type": "function", "function": {"name": "noop"}}],
|
||||
runtime=make_mock_runtime(),
|
||||
state=state,
|
||||
)
|
||||
|
||||
def _make_middleware(self, counter: Any) -> SummarizationMiddleware: # noqa: ANN401
|
||||
"""Middleware whose triggers never fire, so only `_count_tokens` runs."""
|
||||
return SummarizationMiddleware(
|
||||
model=make_mock_model(),
|
||||
backend=MockBackend(),
|
||||
trigger=("tokens", 1_000_000),
|
||||
token_counter=counter,
|
||||
truncate_args_settings={"trigger": ("tokens", 1_000_000), "keep": ("messages", 2)},
|
||||
)
|
||||
|
||||
def test_in_body_typeerror_propagates_sync(self) -> None:
|
||||
# Accepts `tools` but its body raises only when tools is supplied, mimicking
|
||||
# a counter whose tool-schema handling is broken. The old probe swallowed
|
||||
# this and silently recounted without tools; it must now surface loudly.
|
||||
def broken_counter(_messages: list[BaseMessage], *, tools: Any = None) -> int: # noqa: ANN401
|
||||
if tools is not None:
|
||||
msg = "tool schema conversion failed"
|
||||
raise TypeError(msg)
|
||||
return 10
|
||||
|
||||
middleware = self._make_middleware(broken_counter)
|
||||
|
||||
def handler(_req: ModelRequest) -> ModelResponse:
|
||||
return AIMessage(content="ok")
|
||||
|
||||
with pytest.raises(TypeError, match="tool schema conversion failed"):
|
||||
middleware.wrap_model_call(self._make_request_with_tools(), handler)
|
||||
|
||||
async def test_in_body_typeerror_propagates_async(self) -> None:
|
||||
# `awrap_model_call` delegates to the synchronous `_count_tokens`, so the
|
||||
# counting logic is shared with the sync test above. This guards the
|
||||
# async wrapper itself from swallowing the propagated `TypeError`.
|
||||
def broken_counter(_messages: list[BaseMessage], *, tools: Any = None) -> int: # noqa: ANN401
|
||||
if tools is not None:
|
||||
msg = "tool schema conversion failed"
|
||||
raise TypeError(msg)
|
||||
return 10
|
||||
|
||||
middleware = self._make_middleware(broken_counter)
|
||||
|
||||
async def handler(_req: ModelRequest) -> ModelResponse:
|
||||
return AIMessage(content="ok")
|
||||
|
||||
with pytest.raises(TypeError, match="tool schema conversion failed"):
|
||||
await middleware.awrap_model_call(self._make_request_with_tools(), handler)
|
||||
|
||||
def test_counter_without_tools_param_is_called_without_tools(self) -> None:
|
||||
# Tools are present on the request, but a counter that does not declare
|
||||
# `tools` is invoked without them -- and does not raise.
|
||||
def counter(_messages: list[BaseMessage]) -> int:
|
||||
return 10
|
||||
|
||||
middleware = self._make_middleware(counter)
|
||||
captured: ModelRequest | None = None
|
||||
|
||||
def handler(req: ModelRequest) -> ModelResponse:
|
||||
nonlocal captured
|
||||
captured = req
|
||||
return AIMessage(content="ok")
|
||||
|
||||
result = middleware.wrap_model_call(self._make_request_with_tools(), handler)
|
||||
|
||||
assert isinstance(result, AIMessage)
|
||||
assert captured is not None # Handler ran; counting without tools did not raise.
|
||||
|
||||
def test_uninspectable_counter_probes_with_tools(self) -> None:
|
||||
# Signature can't be read, so `_counter_accepts_tools` is `None` and
|
||||
# `_count_tokens` probes with `tools=`. The probe succeeds, so the
|
||||
# tools-inclusive count is returned -- tools are not silently dropped.
|
||||
counter = _OpaqueCounter(lambda _messages, *, tools=None: 99 if tools is not None else 7)
|
||||
middleware = self._make_middleware(counter)
|
||||
assert middleware._counter_accepts_tools is None
|
||||
|
||||
messages = make_conversation_messages()
|
||||
tools = [{"type": "function", "function": {"name": "noop"}}]
|
||||
assert middleware._count_tokens(messages, None, tools) == 99
|
||||
|
||||
def test_uninspectable_counter_falls_back_without_tools(self) -> None:
|
||||
# The opaque counter rejects `tools=` at call time. This is the sole path
|
||||
# that still catches a `TypeError` from the probe, recounting without
|
||||
# tools rather than surfacing it.
|
||||
counter = _OpaqueCounter(lambda _messages: 7)
|
||||
middleware = self._make_middleware(counter)
|
||||
assert middleware._counter_accepts_tools is None
|
||||
|
||||
messages = make_conversation_messages()
|
||||
tools = [{"type": "function", "function": {"name": "noop"}}]
|
||||
assert middleware._count_tokens(messages, None, tools) == 7
|
||||
|
||||
Reference in New Issue
Block a user