core: fix trim_messages token_counter detection for lambdas, subclasses, and string annotations

libs/core/langchain_core/messages/utils.py

@@ -1076,6 +1076,72 @@
     return merged
 
 
+def _is_per_message_token_counter(fn: Callable) -> bool:
+    """Detect whether a callable token counter operates on individual messages.
+
+    Returns ``True`` if *fn* is a **per-message** counter (takes a single
+    :class:`BaseMessage`), ``False`` if it is a **per-list** counter (takes a
+    list of messages).
+
+    The previous implementation used an identity check
+    (``annotation is BaseMessage``) which failed for lambdas, un-annotated
+    functions, string annotations, ``BaseMessage`` sub-class annotations, and
+    postponed annotations (``from __future__ import annotations``).
+
+    The new implementation handles these cases with a three-tier strategy:
+
+    1. **Class annotation** – ``issubclass`` instead of ``is``.
+    2. **String annotation** – match against known message class names.
+    3. **No annotation** – probe by calling the function with a single dummy
+       message and catching ``TypeError`` / ``AttributeError``.
+    """
+    try:
+        sig = inspect.signature(fn)
+    except (ValueError, TypeError):
+        return False
+
+    params = list(sig.parameters.values())
+    if not params:
+        return False
+
+    annotation = params[0].annotation
+
+    # 1. Class annotation: use issubclass to handle BaseMessage sub-classes.
+    if isinstance(annotation, type) and issubclass(annotation, BaseMessage):
+        return True
+
+    # 2. String annotation: covers `from __future__ import annotations` and
+    #    explicit `"BaseMessage"` / `"HumanMessage"` etc.
+    if isinstance(annotation, str):
+        _message_class_names = {
+            cls.__name__
+            for cls in (
+                BaseMessage,
+                HumanMessage,
+                AIMessage,
+                SystemMessage,
+                ChatMessage,
+                FunctionMessage,
+                ToolMessage,
+            )
+        }
+        return annotation in _message_class_names
+
+    # 3. No annotation (lambdas, un-annotated functions): probe by calling with
+    #    a single dummy message.  Built-in list-level callables like ``len``
+    #    will raise ``TypeError`` here because ``BaseMessage`` has no
+    #    ``__len__``; list-iterating counters will raise ``TypeError`` because
+    #    ``BaseMessage`` is not iterable.
+    if annotation is inspect.Parameter.empty:
+        try:
+            fn(HumanMessage(content="a"))
+            return True
+        except (TypeError, AttributeError):
+            return False
+
+    return False
+
+
 # TODO: Update so validation errors (for token_counter, for example) are raised on
 # init not at runtime.
 @_runnable_support
@@ -1417,12 +1483,7 @@
     if hasattr(actual_token_counter, "get_num_tokens_from_messages"):
         list_token_counter = actual_token_counter.get_num_tokens_from_messages
     elif callable(actual_token_counter):
-        if (
-            next(
-                iter(inspect.signature(actual_token_counter).parameters.values())
-            ).annotation
-            is BaseMessage
-        ):
+        if _is_per_message_token_counter(actual_token_counter):
 
             def list_token_counter(messages: Sequence[BaseMessage]) -> int:
                 return sum(actual_token_counter(msg) for msg in messages)  # type: ignore[arg-type, misc]

libs/core/tests/unit_tests/messages/test_utils.py

@@ -2958,3 +2958,81 @@ def get_time(timezone: str) -> str:
     # Test with empty tools list should equal base count
     count_empty_tools = count_tokens_approximately(messages, tools=[])
     assert count_empty_tools == base_count
+
+
+# --- Tests for per-message token_counter detection (issue #35629) ---
+
+_PER_MSG_TRIM_INPUT = [
+    HumanMessage("What is 2 + 2?"),
+    AIMessage("It is 4."),
+    HumanMessage("What about 3 + 3?"),
+]
+
+
+def test_trim_messages_lambda_token_counter() -> None:
+    """Lambda per-message counters have no annotation and should be detected."""
+    result = trim_messages(
+        _PER_MSG_TRIM_INPUT,
+        max_tokens=5,
+        token_counter=lambda msg: len(msg.content.split()),
+        strategy="last",
+    )
+    # "What about 3 + 3?" = 5 tokens, fits exactly
+    assert result == [HumanMessage("What about 3 + 3?")]
+
+
+def test_trim_messages_unannotated_token_counter() -> None:
+    """Un-annotated per-message counters should be detected via probing."""
+
+    def counter_no_annotation(msg):  # noqa: ANN001
+        return len(msg.content.split())
+
+    result = trim_messages(
+        _PER_MSG_TRIM_INPUT,
+        max_tokens=5,
+        token_counter=counter_no_annotation,
+        strategy="last",
+    )
+    assert result == [HumanMessage("What about 3 + 3?")]
+
+
+def test_trim_messages_string_annotation_token_counter() -> None:
+    """String-annotated per-message counters should be detected."""
+
+    def counter_string(msg: "BaseMessage") -> int:  # noqa: UP037
+        return len(msg.content.split())
+
+    result = trim_messages(
+        _PER_MSG_TRIM_INPUT,
+        max_tokens=5,
+        token_counter=counter_string,
+        strategy="last",
+    )
+    assert result == [HumanMessage("What about 3 + 3?")]
+
+
+def test_trim_messages_subclass_annotation_token_counter() -> None:
+    """Counters annotated with a BaseMessage subclass should be detected."""
+
+    def counter_subclass(msg: HumanMessage) -> int:
+        return len(msg.content.split())
+
+    result = trim_messages(
+        _PER_MSG_TRIM_INPUT,
+        max_tokens=5,
+        token_counter=counter_subclass,
+        strategy="last",
+    )
+    assert result == [HumanMessage("What about 3 + 3?")]
+
+
+def test_trim_messages_list_counter_len_still_works() -> None:
+    """Built-in ``len`` (per-list counter) must not be misclassified."""
+    result = trim_messages(
+        _PER_MSG_TRIM_INPUT,
+        max_tokens=2,
+        token_counter=len,
+        strategy="last",
+    )
+    # len counts messages, not tokens — last 2 messages
+    assert result == [AIMessage("It is 4."), HumanMessage("What about 3 + 3?")]