From 5a2cfb49e045988d290a1c7e3a0c589d6b371694 Mon Sep 17 00:00:00 2001
From: William FH <13333726+hinthornw@users.noreply.github.com>
Date: Tue, 29 Oct 2024 21:27:52 -0700
Subject: [PATCH] Support message trimming on single messages (#27729)

Permit trimming message lists of length 1
---
 libs/core/langchain_core/messages/utils.py    |  5 +--
 .../tests/unit_tests/messages/test_utils.py   | 36 +++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/libs/core/langchain_core/messages/utils.py b/libs/core/langchain_core/messages/utils.py
index 9bf617af17d48..0ea03e40d00d4 100644
--- a/libs/core/langchain_core/messages/utils.py
+++ b/libs/core/langchain_core/messages/utils.py
@@ -1210,13 +1210,14 @@ def _first_max_tokens(
     ] = None,
 ) -> list[BaseMessage]:
     messages = list(messages)
+    if not messages:
+        return messages
     idx = 0
     for i in range(len(messages)):
         if token_counter(messages[:-i] if i else messages) <= max_tokens:
             idx = len(messages) - i
             break
-
-    if idx < len(messages) - 1 and partial_strategy:
+    if partial_strategy and (idx < len(messages) - 1 or idx == 0):
         included_partial = False
         if isinstance(messages[idx].content, list):
             excluded = messages[idx].model_copy(deep=True)
diff --git a/libs/core/tests/unit_tests/messages/test_utils.py b/libs/core/tests/unit_tests/messages/test_utils.py
index 9cd9543dc7c8e..ff251e3632078 100644
--- a/libs/core/tests/unit_tests/messages/test_utils.py
+++ b/libs/core/tests/unit_tests/messages/test_utils.py
@@ -301,6 +301,42 @@ def test_trim_messages_last_40_include_system_allow_partial_start_on_human() ->
     assert _MESSAGES_TO_TRIM == _MESSAGES_TO_TRIM_COPY
 
 
+def test_trim_messages_allow_partial_one_message() -> None:
+    expected = [
+        HumanMessage("Th", id="third"),
+    ]
+
+    actual = trim_messages(
+        [HumanMessage("This is a funky text.", id="third")],
+        max_tokens=2,
+        token_counter=lambda messages: sum(len(m.content) for m in messages),
+        text_splitter=lambda x: list(x),
+        strategy="first",
+        allow_partial=True,
+    )
+
+    assert actual == expected
+    assert _MESSAGES_TO_TRIM == _MESSAGES_TO_TRIM_COPY
+
+
+def test_trim_messages_last_allow_partial_one_message() -> None:
+    expected = [
+        HumanMessage("t.", id="third"),
+    ]
+
+    actual = trim_messages(
+        [HumanMessage("This is a funky text.", id="third")],
+        max_tokens=2,
+        token_counter=lambda messages: sum(len(m.content) for m in messages),
+        text_splitter=lambda x: list(x),
+        strategy="last",
+        allow_partial=True,
+    )
+
+    assert actual == expected
+    assert _MESSAGES_TO_TRIM == _MESSAGES_TO_TRIM_COPY
+
+
 def test_trim_messages_allow_partial_text_splitter() -> None:
     expected = [
         HumanMessage("a 4 token text.", id="third"),