fix(core): robust per-message token counter detection in trim_messages

libs/core/langchain_core/messages/utils.py

@@ -14,7 +14,7 @@
 import json
 import logging
 import math
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import Callable, Iterable, MutableSequence, Sequence
 from functools import partial, wraps
 from typing import (
     TYPE_CHECKING,
@@ -26,6 +26,8 @@
     Protocol,
     TypeVar,
     cast,
+    get_origin,
+    get_type_hints,
     overload,
 )
 from xml.sax.saxutils import escape, quoteattr
@@ -1076,6 +1078,72 @@ def merge_message_runs(
     return merged
 
 
+def _is_per_message_token_counter(fn: Callable) -> bool:
+    """Determine if a callable token counter operates per-message or per-list.
+
+    Uses ``typing.get_type_hints`` to robustly resolve the first parameter's
+    annotation (handles ``from __future__ import annotations``, string
+    annotations, and subclasses of ``BaseMessage``).  Falls back to raw
+    ``inspect.signature`` when ``get_type_hints`` fails.
+
+    Returns ``True`` when the callable should be called once *per message*
+    (i.e. its first parameter is annotated as a ``BaseMessage`` or subclass).
+    Unannotated callables (lambdas, bare ``def``) are treated as per-list for
+    backward compatibility.
+    """
+    annotation = None
+
+    # 1. Try to get the resolved type hint (handles string / postponed annotations)
+    try:
+        hints = get_type_hints(fn)
+        if hints:
+            annotation = next(iter(hints.values()))
+    except Exception:
+        logger.debug("get_type_hints failed for %s", fn)
+
+    # 2. Fall back to raw annotation from inspect.signature
+    if annotation is None:
+        try:
+            param = next(iter(inspect.signature(fn).parameters.values()))
+            if param.annotation is not inspect.Parameter.empty:
+                annotation = param.annotation
+        except (StopIteration, ValueError, TypeError):
+            pass
+
+    # 3. No annotation at all (lambdas, bare functions) → per-list (backward compat)
+    if annotation is None:
+        return False
+
+    # 4. If the annotation is a string we couldn't resolve, check for message-like
+    if isinstance(annotation, str):
+        lowered = annotation.lower()
+        if any(kw in lowered for kw in ("list", "sequence", "iterable")):
+            return False
+        return any(kw in lowered for kw in ("message", "basemessage"))
+
+    # 5. Check if annotation is a generic collection type (list[...], Sequence[...])
+    origin = get_origin(annotation)
+    if origin is not None and origin in (
+        list,
+        tuple,
+        set,
+        frozenset,
+        Sequence,
+        MutableSequence,
+        Iterable,
+    ):
+        return False
+
+    # 6. Check if annotation is a concrete collection class
+    if isinstance(annotation, type) and issubclass(
+        annotation, (list, tuple, set, frozenset, Sequence)
+    ):
+        return False
+
+    # 7. Check if annotation is BaseMessage or a subclass → per-message
+    return isinstance(annotation, type) and issubclass(annotation, BaseMessage)
+
+
 # TODO: Update so validation errors (for token_counter, for example) are raised on
 # init not at runtime.
 @_runnable_support
@@ -1417,12 +1485,7 @@ def dummy_token_counter(messages: list[BaseMessage]) -> int:
     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

@@ -753,6 +753,40 @@ def test_trim_messages_token_counter_shortcut_with_options() -> None:
     assert messages == messages_copy
 
 
+def test_trim_messages_per_message_counter_with_subclass_annotation() -> None:
+    """Per-message counter annotated with a BaseMessage subclass is detected."""
+    messages = [
+        HumanMessage("What is 2 + 2?"),
+        AIMessage("It is 4."),
+        HumanMessage("What about 3 + 3?"),
+    ]
+
+    def counter(msg: HumanMessage) -> int:
+        return len(msg.content.split())
+
+    result = trim_messages(messages, max_tokens=5, token_counter=counter)
+    assert len(result) >= 1
+    assert all(isinstance(m, BaseMessage) for m in result)
+
+
+def test_trim_messages_per_message_counter_with_postponed_annotations() -> None:
+    """Per-message counter with string (postponed) annotation is detected."""
+    messages = [
+        HumanMessage("What is 2 + 2?"),
+        AIMessage("It is 4."),
+        HumanMessage("What about 3 + 3?"),
+    ]
+
+    # Simulates `from __future__ import annotations` effect:
+    # get_type_hints resolves the string to the actual class
+    def counter(msg: BaseMessage) -> int:
+        return len(msg.content.split())
+
+    result = trim_messages(messages, max_tokens=5, token_counter=counter)
+    assert len(result) >= 1
+    assert all(isinstance(m, BaseMessage) for m in result)
+
+
 class FakeTokenCountingModel(FakeChatModel):
     @override
     def get_num_tokens_from_messages(