From da59eb7eb4d97ab50a2cafff0c770e089be7d77a Mon Sep 17 00:00:00 2001 From: Ben Gladwell Date: Wed, 30 Apr 2025 17:56:22 -0400 Subject: [PATCH] anthropic: Allow kwargs to pass through when counting tokens (#31082) - **Description:** `ChatAnthropic.get_num_tokens_from_messages` does not currently receive `kwargs` and pass those on to `self._client.beta.messages.count_tokens`. This is a problem if you need to pass specific options to `count_tokens`, such as the `thinking` option. This PR fixes that. - **Issue:** N/A - **Dependencies:** None - **Twitter handle:** @bengladwell Co-authored-by: ccurme --- .../anthropic/langchain_anthropic/chat_models.py | 2 +- .../anthropic/tests/unit_tests/test_chat_models.py | 14 ++++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/libs/partners/anthropic/langchain_anthropic/chat_models.py b/libs/partners/anthropic/langchain_anthropic/chat_models.py index f872fe32c18..d1f0237b9b9 100644 --- a/libs/partners/anthropic/langchain_anthropic/chat_models.py +++ b/libs/partners/anthropic/langchain_anthropic/chat_models.py @@ -1588,6 +1588,7 @@ class ChatAnthropic(BaseChatModel): tools: Optional[ Sequence[Union[dict[str, Any], type, Callable, BaseTool]] ] = None, + **kwargs: Any, ) -> int: """Count tokens in a sequence of input messages. @@ -1647,7 +1648,6 @@ class ChatAnthropic(BaseChatModel): https://docs.anthropic.com/en/docs/build-with-claude/token-counting """ formatted_system, formatted_messages = _format_messages(messages) - kwargs: dict[str, Any] = {} if isinstance(formatted_system, str): kwargs["system"] = formatted_system if tools: diff --git a/libs/partners/anthropic/tests/unit_tests/test_chat_models.py b/libs/partners/anthropic/tests/unit_tests/test_chat_models.py index 008b2be1cb6..4d50794b00b 100644 --- a/libs/partners/anthropic/tests/unit_tests/test_chat_models.py +++ b/libs/partners/anthropic/tests/unit_tests/test_chat_models.py @@ -2,7 +2,9 @@ import os from typing import Any, Callable, Literal, cast +from unittest.mock import patch +import anthropic import pytest from anthropic.types import Message, TextBlock, Usage from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage @@ -940,3 +942,15 @@ def test_optional_description() -> None: sample_field: str _ = llm.with_structured_output(SampleModel.model_json_schema()) + + +def test_get_num_tokens_from_messages_passes_kwargs() -> None: + """Test that get_num_tokens_from_messages passes kwargs to the model.""" + llm = ChatAnthropic(model="claude-3-5-haiku-latest") + + with patch.object(anthropic, "Client") as _Client: + llm.get_num_tokens_from_messages([HumanMessage("foo")], foo="bar") + + assert ( + _Client.return_value.beta.messages.count_tokens.call_args.kwargs["foo"] == "bar" + )