mirror of
https://github.com/hwchase17/langchain.git
synced 2025-07-11 23:40:24 +00:00
Support message trimming on single messages
This commit is contained in:
parent
0496c15123
commit
f7078e08da
@ -876,13 +876,14 @@ def _first_max_tokens(
|
|||||||
] = None,
|
] = None,
|
||||||
) -> list[BaseMessage]:
|
) -> list[BaseMessage]:
|
||||||
messages = list(messages)
|
messages = list(messages)
|
||||||
|
if not messages:
|
||||||
|
return messages
|
||||||
idx = 0
|
idx = 0
|
||||||
for i in range(len(messages)):
|
for i in range(len(messages)):
|
||||||
if token_counter(messages[:-i] if i else messages) <= max_tokens:
|
if token_counter(messages[:-i] if i else messages) <= max_tokens:
|
||||||
idx = len(messages) - i
|
idx = len(messages) - i
|
||||||
break
|
break
|
||||||
|
if partial_strategy and (idx < len(messages) - 1 or idx == 0):
|
||||||
if idx < len(messages) - 1 and partial_strategy:
|
|
||||||
included_partial = False
|
included_partial = False
|
||||||
if isinstance(messages[idx].content, list):
|
if isinstance(messages[idx].content, list):
|
||||||
excluded = messages[idx].model_copy(deep=True)
|
excluded = messages[idx].model_copy(deep=True)
|
||||||
|
@ -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
|
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:
|
def test_trim_messages_allow_partial_text_splitter() -> None:
|
||||||
expected = [
|
expected = [
|
||||||
HumanMessage("a 4 token text.", id="third"),
|
HumanMessage("a 4 token text.", id="third"),
|
||||||
|
Loading…
Reference in New Issue
Block a user