mirror of
https://github.com/langchain-ai/langchain.git
synced 2026-10-05 09:25:14 +03:00
feat(core): add 'approximate' alias in place of count_tokens_approximately (#33045)
### Description: earlier we have to use like below: ```python from langchain_core.messages import trim_messages from langchain_core.messages.utils import count_tokens_approximately trim_messages(..., token_counter=count_tokens_approximately) ``` Now can be used as like this also ```python from langchain_core.messages import trim_messages trim_messages(..., token_counter="approximate") ``` - [x] **Added tests** - [x] **Lint and test**: Run this as I made change in langchain/core, uv run --group test pytest tests/unit_tests/messages/test_utils.py -v <img width="1006" height="66" alt="image" src="https://github.com/user-attachments/assets/c6938c29-a781-4e7f-871b-8e888ee764b7" /> --------- Co-authored-by: Mason Daugherty <mason@langchain.dev> Co-authored-by: Mason Daugherty <github@mdrxy.com>
This commit is contained in:
2 files changed
+192
-46
No files matched your search
@@ -720,7 +720,8 @@ def trim_messages(
|
|||||||
max_tokens: int,
|
max_tokens: int,
|
||||||
token_counter: Callable[[list[BaseMessage]], int]
|
token_counter: Callable[[list[BaseMessage]], int]
|
||||||
| Callable[[BaseMessage], int]
|
| Callable[[BaseMessage], int]
|
||||||
| BaseLanguageModel,
|
| BaseLanguageModel
|
||||||
|
| Literal["approximate"],
|
||||||
strategy: Literal["first", "last"] = "last",
|
strategy: Literal["first", "last"] = "last",
|
||||||
allow_partial: bool = False,
|
allow_partial: bool = False,
|
||||||
end_on: str | type[BaseMessage] | Sequence[str | type[BaseMessage]] | None = None,
|
end_on: str | type[BaseMessage] | Sequence[str | type[BaseMessage]] | None = None,
|
||||||
@@ -758,53 +759,65 @@ def trim_messages(
|
|||||||
messages: Sequence of Message-like objects to trim.
|
messages: Sequence of Message-like objects to trim.
|
||||||
max_tokens: Max token count of trimmed messages.
|
max_tokens: Max token count of trimmed messages.
|
||||||
token_counter: Function or llm for counting tokens in a `BaseMessage` or a
|
token_counter: Function or llm for counting tokens in a `BaseMessage` or a
|
||||||
list of `BaseMessage`. If a `BaseLanguageModel` is passed in then
|
list of `BaseMessage`.
|
||||||
`BaseLanguageModel.get_num_tokens_from_messages()` will be used.
|
|
||||||
Set to `len` to count the number of **messages** in the chat history.
|
If a `BaseLanguageModel` is passed in then
|
||||||
|
`BaseLanguageModel.get_num_tokens_from_messages()` will be used. Set to
|
||||||
|
`len` to count the number of **messages** in the chat history.
|
||||||
|
|
||||||
|
You can also use string shortcuts for convenience:
|
||||||
|
|
||||||
|
- `'approximate'`: Uses `count_tokens_approximately` for fast, approximate
|
||||||
|
token counts.
|
||||||
|
|
||||||
!!! note
|
!!! note
|
||||||
|
|
||||||
Use `count_tokens_approximately` to get fast, approximate token
|
`count_tokens_approximately` (or the shortcut `'approximate'`) is
|
||||||
counts.
|
recommended for using `trim_messages` on the hot path, where exact token
|
||||||
|
counting is not necessary.
|
||||||
This is recommended for using `trim_messages` on the hot path, where
|
|
||||||
exact token counting is not necessary.
|
|
||||||
|
|
||||||
strategy: Strategy for trimming.
|
strategy: Strategy for trimming.
|
||||||
|
|
||||||
- `'first'`: Keep the first `<= n_count` tokens of the messages.
|
- `'first'`: Keep the first `<= n_count` tokens of the messages.
|
||||||
- `'last'`: Keep the last `<= n_count` tokens of the messages.
|
- `'last'`: Keep the last `<= n_count` tokens of the messages.
|
||||||
allow_partial: Whether to split a message if only part of the message can be
|
allow_partial: Whether to split a message if only part of the message can be
|
||||||
included. If `strategy='last'` then the last partial contents of a message
|
included.
|
||||||
are included. If `strategy='first'` then the first partial contents of a
|
|
||||||
message are included.
|
|
||||||
end_on: The message type to end on. If specified then every message after the
|
|
||||||
last occurrence of this type is ignored. If `strategy='last'` then this
|
|
||||||
is done before we attempt to get the last `max_tokens`. If
|
|
||||||
`strategy='first'` then this is done after we get the first
|
|
||||||
`max_tokens`. Can be specified as string names (e.g. `'system'`,
|
|
||||||
`'human'`, `'ai'`, ...) or as `BaseMessage` classes (e.g.
|
|
||||||
`SystemMessage`, `HumanMessage`, `AIMessage`, ...). Can be a single
|
|
||||||
type or a list of types.
|
|
||||||
|
|
||||||
start_on: The message type to start on. Should only be specified if
|
If `strategy='last'` then the last partial contents of a message are
|
||||||
`strategy='last'`. If specified then every message before
|
included. If `strategy='first'` then the first partial contents of a
|
||||||
the first occurrence of this type is ignored. This is done after we trim
|
message are included.
|
||||||
the initial messages to the last `max_tokens`. Does not
|
end_on: The message type to end on.
|
||||||
apply to a `SystemMessage` at index 0 if `include_system=True`. Can be
|
|
||||||
specified as string names (e.g. `'system'`, `'human'`, `'ai'`, ...) or
|
If specified then every message after the last occurrence of this type is
|
||||||
as `BaseMessage` classes (e.g. `SystemMessage`, `HumanMessage`,
|
ignored. If `strategy='last'` then this is done before we attempt to get the
|
||||||
`AIMessage`, ...). Can be a single type or a list of types.
|
last `max_tokens`. If `strategy='first'` then this is done after we get the
|
||||||
|
first `max_tokens`. Can be specified as string names (e.g. `'system'`,
|
||||||
|
`'human'`, `'ai'`, ...) or as `BaseMessage` classes (e.g. `SystemMessage`,
|
||||||
|
`HumanMessage`, `AIMessage`, ...). Can be a single type or a list of types.
|
||||||
|
|
||||||
|
start_on: The message type to start on.
|
||||||
|
|
||||||
|
Should only be specified if `strategy='last'`. If specified then every
|
||||||
|
message before the first occurrence of this type is ignored. This is done
|
||||||
|
after we trim the initial messages to the last `max_tokens`. Does not apply
|
||||||
|
to a `SystemMessage` at index 0 if `include_system=True`. Can be specified
|
||||||
|
as string names (e.g. `'system'`, `'human'`, `'ai'`, ...) or as
|
||||||
|
`BaseMessage` classes (e.g. `SystemMessage`, `HumanMessage`, `AIMessage`,
|
||||||
|
...). Can be a single type or a list of types.
|
||||||
|
|
||||||
include_system: Whether to keep the `SystemMessage` if there is one at index
|
include_system: Whether to keep the `SystemMessage` if there is one at index
|
||||||
`0`. Should only be specified if `strategy="last"`.
|
`0`.
|
||||||
|
|
||||||
|
Should only be specified if `strategy="last"`.
|
||||||
text_splitter: Function or `langchain_text_splitters.TextSplitter` for
|
text_splitter: Function or `langchain_text_splitters.TextSplitter` for
|
||||||
splitting the string contents of a message. Only used if
|
splitting the string contents of a message.
|
||||||
`allow_partial=True`. If `strategy='last'` then the last split tokens
|
|
||||||
from a partial message will be included. if `strategy='first'` then the
|
Only used if `allow_partial=True`. If `strategy='last'` then the last split
|
||||||
first split tokens from a partial message will be included. Token splitter
|
tokens from a partial message will be included. if `strategy='first'` then
|
||||||
assumes that separators are kept, so that split contents can be directly
|
the first split tokens from a partial message will be included. Token
|
||||||
concatenated to recreate the original text. Defaults to splitting on
|
splitter assumes that separators are kept, so that split contents can be
|
||||||
newlines.
|
directly concatenated to recreate the original text. Defaults to splitting
|
||||||
|
on newlines.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
List of trimmed `BaseMessage`.
|
List of trimmed `BaseMessage`.
|
||||||
@@ -815,8 +828,8 @@ def trim_messages(
|
|||||||
|
|
||||||
Example:
|
Example:
|
||||||
Trim chat history based on token count, keeping the `SystemMessage` if
|
Trim chat history based on token count, keeping the `SystemMessage` if
|
||||||
present, and ensuring that the chat history starts with a `HumanMessage` (
|
present, and ensuring that the chat history starts with a `HumanMessage` (or a
|
||||||
or a `SystemMessage` followed by a `HumanMessage`).
|
`SystemMessage` followed by a `HumanMessage`).
|
||||||
|
|
||||||
```python
|
```python
|
||||||
from langchain_core.messages import (
|
from langchain_core.messages import (
|
||||||
@@ -869,8 +882,34 @@ def trim_messages(
|
|||||||
]
|
]
|
||||||
```
|
```
|
||||||
|
|
||||||
|
Trim chat history using approximate token counting with `'approximate'`:
|
||||||
|
|
||||||
|
```python
|
||||||
|
trim_messages(
|
||||||
|
messages,
|
||||||
|
max_tokens=45,
|
||||||
|
strategy="last",
|
||||||
|
# Using the "approximate" shortcut for fast token counting
|
||||||
|
token_counter="approximate",
|
||||||
|
start_on="human",
|
||||||
|
include_system=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
# This is equivalent to using `count_tokens_approximately` directly
|
||||||
|
from langchain_core.messages.utils import count_tokens_approximately
|
||||||
|
|
||||||
|
trim_messages(
|
||||||
|
messages,
|
||||||
|
max_tokens=45,
|
||||||
|
strategy="last",
|
||||||
|
token_counter=count_tokens_approximately,
|
||||||
|
start_on="human",
|
||||||
|
include_system=True,
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
Trim chat history based on the message count, keeping the `SystemMessage` if
|
Trim chat history based on the message count, keeping the `SystemMessage` if
|
||||||
present, and ensuring that the chat history starts with a `HumanMessage` (
|
present, and ensuring that the chat history starts with a HumanMessage (
|
||||||
or a `SystemMessage` followed by a `HumanMessage`).
|
or a `SystemMessage` followed by a `HumanMessage`).
|
||||||
|
|
||||||
trim_messages(
|
trim_messages(
|
||||||
@@ -992,24 +1031,44 @@ def trim_messages(
|
|||||||
raise ValueError(msg)
|
raise ValueError(msg)
|
||||||
|
|
||||||
messages = convert_to_messages(messages)
|
messages = convert_to_messages(messages)
|
||||||
if hasattr(token_counter, "get_num_tokens_from_messages"):
|
|
||||||
list_token_counter = token_counter.get_num_tokens_from_messages
|
# Handle string shortcuts for token counter
|
||||||
elif callable(token_counter):
|
if isinstance(token_counter, str):
|
||||||
|
if token_counter in _TOKEN_COUNTER_SHORTCUTS:
|
||||||
|
actual_token_counter = _TOKEN_COUNTER_SHORTCUTS[token_counter]
|
||||||
|
else:
|
||||||
|
available_shortcuts = ", ".join(
|
||||||
|
f"'{key}'" for key in _TOKEN_COUNTER_SHORTCUTS
|
||||||
|
)
|
||||||
|
msg = (
|
||||||
|
f"Invalid token_counter shortcut '{token_counter}'. "
|
||||||
|
f"Available shortcuts: {available_shortcuts}."
|
||||||
|
)
|
||||||
|
raise ValueError(msg)
|
||||||
|
else:
|
||||||
|
# Type narrowing: at this point token_counter is not a str
|
||||||
|
actual_token_counter = token_counter # type: ignore[assignment]
|
||||||
|
|
||||||
|
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 (
|
if (
|
||||||
next(iter(inspect.signature(token_counter).parameters.values())).annotation
|
next(
|
||||||
|
iter(inspect.signature(actual_token_counter).parameters.values())
|
||||||
|
).annotation
|
||||||
is BaseMessage
|
is BaseMessage
|
||||||
):
|
):
|
||||||
|
|
||||||
def list_token_counter(messages: Sequence[BaseMessage]) -> int:
|
def list_token_counter(messages: Sequence[BaseMessage]) -> int:
|
||||||
return sum(token_counter(msg) for msg in messages) # type: ignore[arg-type, misc]
|
return sum(actual_token_counter(msg) for msg in messages) # type: ignore[arg-type, misc]
|
||||||
|
|
||||||
else:
|
else:
|
||||||
list_token_counter = token_counter
|
list_token_counter = actual_token_counter
|
||||||
else:
|
else:
|
||||||
msg = (
|
msg = (
|
||||||
f"'token_counter' expected to be a model that implements "
|
f"'token_counter' expected to be a model that implements "
|
||||||
f"'get_num_tokens_from_messages()' or a function. Received object of type "
|
f"'get_num_tokens_from_messages()' or a function. Received object of type "
|
||||||
f"{type(token_counter)}."
|
f"{type(actual_token_counter)}."
|
||||||
)
|
)
|
||||||
raise ValueError(msg)
|
raise ValueError(msg)
|
||||||
|
|
||||||
@@ -1807,3 +1866,14 @@ def count_tokens_approximately(
|
|||||||
|
|
||||||
# round up once more time in case extra_tokens_per_message is a float
|
# round up once more time in case extra_tokens_per_message is a float
|
||||||
return math.ceil(token_count)
|
return math.ceil(token_count)
|
||||||
|
|
||||||
|
|
||||||
|
# Mapping from string shortcuts to token counter functions
|
||||||
|
def _approximate_token_counter(messages: Sequence[BaseMessage]) -> int:
|
||||||
|
"""Wrapper for `count_tokens_approximately` that matches expected signature."""
|
||||||
|
return count_tokens_approximately(messages)
|
||||||
|
|
||||||
|
|
||||||
|
_TOKEN_COUNTER_SHORTCUTS = {
|
||||||
|
"approximate": _approximate_token_counter,
|
||||||
|
}
|
||||||
@@ -673,6 +673,82 @@ def test_trim_messages_start_on_with_allow_partial() -> None:
|
|||||||
assert messages == messages_copy
|
assert messages == messages_copy
|
||||||
|
|
||||||
|
|
||||||
|
def test_trim_messages_token_counter_shortcut_approximate() -> None:
|
||||||
|
"""Test that `'approximate'` shortcut works for `token_counter`."""
|
||||||
|
messages = [
|
||||||
|
SystemMessage("This is a test message"),
|
||||||
|
HumanMessage("Another test message", id="first"),
|
||||||
|
AIMessage("AI response here", id="second"),
|
||||||
|
]
|
||||||
|
messages_copy = [m.model_copy(deep=True) for m in messages]
|
||||||
|
|
||||||
|
# Test using the "approximate" shortcut
|
||||||
|
result_shortcut = trim_messages(
|
||||||
|
messages,
|
||||||
|
max_tokens=50,
|
||||||
|
token_counter="approximate",
|
||||||
|
strategy="last",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Test using count_tokens_approximately directly
|
||||||
|
result_direct = trim_messages(
|
||||||
|
messages,
|
||||||
|
max_tokens=50,
|
||||||
|
token_counter=count_tokens_approximately,
|
||||||
|
strategy="last",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Both should produce the same result
|
||||||
|
assert result_shortcut == result_direct
|
||||||
|
assert messages == messages_copy
|
||||||
|
|
||||||
|
|
||||||
|
def test_trim_messages_token_counter_shortcut_invalid() -> None:
|
||||||
|
"""Test that invalid `token_counter` shortcut raises `ValueError`."""
|
||||||
|
messages = [
|
||||||
|
SystemMessage("This is a test message"),
|
||||||
|
HumanMessage("Another test message"),
|
||||||
|
]
|
||||||
|
|
||||||
|
# Test with invalid shortcut - intentionally passing invalid string to verify
|
||||||
|
# runtime error handling for dynamically-constructed inputs
|
||||||
|
with pytest.raises(ValueError, match="Invalid token_counter shortcut 'invalid'"):
|
||||||
|
trim_messages( # type: ignore[call-overload]
|
||||||
|
messages,
|
||||||
|
max_tokens=50,
|
||||||
|
token_counter="invalid",
|
||||||
|
strategy="last",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_trim_messages_token_counter_shortcut_with_options() -> None:
|
||||||
|
"""Test that `'approximate'` shortcut works with different trim options."""
|
||||||
|
messages = [
|
||||||
|
SystemMessage("System instructions"),
|
||||||
|
HumanMessage("First human message", id="first"),
|
||||||
|
AIMessage("First AI response", id="ai1"),
|
||||||
|
HumanMessage("Second human message", id="second"),
|
||||||
|
AIMessage("Second AI response", id="ai2"),
|
||||||
|
]
|
||||||
|
messages_copy = [m.model_copy(deep=True) for m in messages]
|
||||||
|
|
||||||
|
# Test with various options
|
||||||
|
result = trim_messages(
|
||||||
|
messages,
|
||||||
|
max_tokens=100,
|
||||||
|
token_counter="approximate",
|
||||||
|
strategy="last",
|
||||||
|
include_system=True,
|
||||||
|
start_on="human",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Should include system message and start on human
|
||||||
|
assert len(result) >= 2
|
||||||
|
assert isinstance(result[0], SystemMessage)
|
||||||
|
assert any(isinstance(msg, HumanMessage) for msg in result[1:])
|
||||||
|
assert messages == messages_copy
|
||||||
|
|
||||||
|
|
||||||
class FakeTokenCountingModel(FakeChatModel):
|
class FakeTokenCountingModel(FakeChatModel):
|
||||||
@override
|
@override
|
||||||
def get_num_tokens_from_messages(
|
def get_num_tokens_from_messages(
|
||||||
|
|||||||
Reference in new issue
Block a user