Skip to content

Commit 5a2cfb4

Browse files
authored
Support message trimming on single messages (#27729)
Permit trimming message lists of length 1
1 parent 5111063 commit 5a2cfb4

File tree

2 files changed

+39
-2
lines changed

2 files changed

+39
-2
lines changed

libs/core/langchain_core/messages/utils.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1210,13 +1210,14 @@ def _first_max_tokens(
12101210
] = None,
12111211
) -> list[BaseMessage]:
12121212
messages = list(messages)
1213+
if not messages:
1214+
return messages
12131215
idx = 0
12141216
for i in range(len(messages)):
12151217
if token_counter(messages[:-i] if i else messages) <= max_tokens:
12161218
idx = len(messages) - i
12171219
break
1218-
1219-
if idx < len(messages) - 1 and partial_strategy:
1220+
if partial_strategy and (idx < len(messages) - 1 or idx == 0):
12201221
included_partial = False
12211222
if isinstance(messages[idx].content, list):
12221223
excluded = messages[idx].model_copy(deep=True)

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,6 +301,42 @@ def test_trim_messages_last_40_include_system_allow_partial_start_on_human() ->
301301
assert _MESSAGES_TO_TRIM == _MESSAGES_TO_TRIM_COPY
302302

303303

304+
def test_trim_messages_allow_partial_one_message() -> None:
305+
expected = [
306+
HumanMessage("Th", id="third"),
307+
]
308+
309+
actual = trim_messages(
310+
[HumanMessage("This is a funky text.", id="third")],
311+
max_tokens=2,
312+
token_counter=lambda messages: sum(len(m.content) for m in messages),
313+
text_splitter=lambda x: list(x),
314+
strategy="first",
315+
allow_partial=True,
316+
)
317+
318+
assert actual == expected
319+
assert _MESSAGES_TO_TRIM == _MESSAGES_TO_TRIM_COPY
320+
321+
322+
def test_trim_messages_last_allow_partial_one_message() -> None:
323+
expected = [
324+
HumanMessage("t.", id="third"),
325+
]
326+
327+
actual = trim_messages(
328+
[HumanMessage("This is a funky text.", id="third")],
329+
max_tokens=2,
330+
token_counter=lambda messages: sum(len(m.content) for m in messages),
331+
text_splitter=lambda x: list(x),
332+
strategy="last",
333+
allow_partial=True,
334+
)
335+
336+
assert actual == expected
337+
assert _MESSAGES_TO_TRIM == _MESSAGES_TO_TRIM_COPY
338+
339+
304340
def test_trim_messages_allow_partial_text_splitter() -> None:
305341
expected = [
306342
HumanMessage("a 4 token text.", id="third"),

0 commit comments

Comments
 (0)