Support message trimming on single messages

This commit is contained in:
William Fu-Hinthorn 2024-10-29 18:51:23 -07:00
parent 0496c15123
commit f7078e08da
2 changed files with 21 additions and 2 deletions

View File

@ -876,13 +876,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)

View File

@ -299,6 +299,24 @@ 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 4 token 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_allow_partial_text_splitter() -> None:
expected = [
HumanMessage("a 4 token text.", id="third"),