From 5ea2c21d5e7ab555366f16c69d9c9bd34615c5fc Mon Sep 17 00:00:00 2001 From: Michael Chin Date: Thu, 16 Oct 2025 20:41:04 -0700 Subject: [PATCH 1/3] fix(aws): Handle thinking with tools, use CountTokens API for Haiku 4.5 --- libs/aws/langchain_aws/chat_models/bedrock.py | 7 ++++++- .../chat_models/bedrock_converse.py | 18 ++++++++++++++++-- libs/aws/langchain_aws/utils.py | 8 +++++++- .../chat_models/test_bedrock_converse.py | 7 +++++++ .../chat_models/test_bedrock_converse.py | 1 + libs/aws/tests/unit_tests/test_utils.py | 1 + 6 files changed, 38 insertions(+), 4 deletions(-) diff --git a/libs/aws/langchain_aws/chat_models/bedrock.py b/libs/aws/langchain_aws/chat_models/bedrock.py index 186b2b7b..43649a73 100644 --- a/libs/aws/langchain_aws/chat_models/bedrock.py +++ b/libs/aws/langchain_aws/chat_models/bedrock.py @@ -1231,7 +1231,12 @@ def bind_tools( base_model = self._get_base_model() if any( x in base_model - for x in ("claude-3-7-", "claude-opus-4-", "claude-sonnet-4-") + for x in ( + "claude-3-7-", + "claude-opus-4-", + "claude-sonnet-4-", + "claude-haiku-4-", + ) ) and thinking_in_params(self.model_kwargs or {}): forced = False if isinstance(tool_choice, bool): diff --git a/libs/aws/langchain_aws/chat_models/bedrock_converse.py b/libs/aws/langchain_aws/chat_models/bedrock_converse.py index 6b461647..7333d7a5 100644 --- a/libs/aws/langchain_aws/chat_models/bedrock_converse.py +++ b/libs/aws/langchain_aws/chat_models/bedrock_converse.py @@ -759,6 +759,7 @@ def validate_environment(self) -> Self: "claude-3-7-sonnet", "claude-sonnet-4", "claude-opus-4", + "claude-haiku-4", ) thinking_params = (self.additional_model_request_fields or {}).get( "thinking", {} @@ -971,20 +972,29 @@ def _get_llm_for_structured_output_no_tool_choice( self, schema: Union[Dict, type], ) -> Runnable[LanguageModelInput, BaseMessage]: + print("In _get_llm_for_structured_output_no_tool_choice") admonition = ( "ChatBedrockConverse structured output relies on forced tool calling, " "which is not supported for this model. This method will raise " "langchain_core.exceptions.OutputParserException if tool calls are not " "generated. Consider adjusting your prompt to ensure the tool is called." ) - if "claude-3-7-sonnet" in self._get_base_model(): + thinking_claude_models = ( + "claude-3-7-sonnet", + "claude-sonnet-4", + "claude-opus-4", + "claude-haiku-4", + ) + if any(model in self._get_base_model() for model in thinking_claude_models): additional_context = ( - "For Claude 3.7 Sonnet models, you can also support forced tool use " + "For Claude 3/4 models, you can also support forced tool use " "by disabling `thinking`." ) admonition = f"{admonition} {additional_context}" + print("Returning warning for thinking claude model") warnings.warn(admonition) try: + print("Attempting bind_tools with structured output") llm = self.bind_tools( [schema], ls_structured_output_format={ @@ -993,11 +1003,14 @@ def _get_llm_for_structured_output_no_tool_choice( }, ) except Exception: + print("Failed, using regular bind_tools") llm = self.bind_tools([schema]) def _raise_if_no_tool_calls(message: AIMessage) -> AIMessage: if not message.tool_calls: + print("No tool calls found, raising OutputParserException") raise OutputParserException(admonition) + print("Tool calls found, returning normally") return message return llm | _raise_if_no_tool_calls @@ -1066,6 +1079,7 @@ def with_structured_output( "claude-3-7-sonnet", "claude-sonnet-4", "claude-opus-4", + "claude-haiku-4", ) if tool_choice is None and any( model in self._get_base_model() for model in thinking_claude_models diff --git a/libs/aws/langchain_aws/utils.py b/libs/aws/langchain_aws/utils.py index c44534b2..355cac58 100644 --- a/libs/aws/langchain_aws/utils.py +++ b/libs/aws/langchain_aws/utils.py @@ -96,7 +96,13 @@ def anthropic_tokens_supported() -> bool: def count_tokens_api_supported_for_model(model: str) -> bool: return any( x in model - for x in ("claude-3-5-", "claude-3-7-", "claude-opus-4-", "claude-sonnet-4-") + for x in ( + "claude-3-5-", + "claude-3-7-", + "claude-opus-4-", + "claude-sonnet-4-", + "claude-haiku-4-", + ) ) diff --git a/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py b/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py index cefc0432..c45355a3 100644 --- a/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py +++ b/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py @@ -431,6 +431,10 @@ def test_guardrails() -> None: "us.anthropic.claude-3-7-sonnet-20250219-v1:0", "us.anthropic.claude-sonnet-4-20250514-v1:0", "us.anthropic.claude-opus-4-20250514-v1:0", + "us.anthropic.claude-opus-4-1-20250805-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + # + #"us.anthropic.claude-haiku-4-5-20251001-v1:0", ], ) def test_structured_output_tool_choice_not_supported(thinking_model: str) -> None: @@ -450,11 +454,14 @@ def test_structured_output_tool_choice_not_supported(thinking_model: str) -> Non "thinking": {"type": "enabled", "budget_tokens": 2000} }, ) + print("Generating structured llm with thinking...") with pytest.warns(match="structured output"): structured_llm = llm.with_structured_output(ClassifyQuery) + print("Invoking on structured llm with thinking and related query...") response = structured_llm.invoke("How big are cats?") assert isinstance(response, ClassifyQuery) + print("Invoking on structured llm with thinking and generic query...") with pytest.raises(OutputParserException): structured_llm.invoke("Hello!") diff --git a/libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py b/libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py index 422434f9..b78c783a 100644 --- a/libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py +++ b/libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py @@ -120,6 +120,7 @@ def test_anthropic_bind_tools_tool_choice() -> None: "anthropic.claude-3-7-sonnet-20250219-v1:0", "anthropic.claude-sonnet-4-20250514-v1:0", "anthropic.claude-opus-4-20250514-v1:0", + "anthropic.claude-haiku-4-5-20251001-v1:0" ], ) def test_anthropic_thinking_bind_tools_tool_choice(thinking_model: str) -> None: diff --git a/libs/aws/tests/unit_tests/test_utils.py b/libs/aws/tests/unit_tests/test_utils.py index 4206ea63..6fa6054b 100644 --- a/libs/aws/tests/unit_tests/test_utils.py +++ b/libs/aws/tests/unit_tests/test_utils.py @@ -435,6 +435,7 @@ def test_trim_message_whitespace_with_empty_messages() -> None: @pytest.mark.parametrize( "model_id,expected_result", [ + ("us.anthropic.claude-haiku-4-5-20251001-v1:0", True), ("us.anthropic.claude-opus-4-20250514-v1:0", True), ("us.anthropic.claude-sonnet-4-20250514-v1:0", True), ("us.anthropic.claude-3-7-sonnet-20250219-v1:0", True), From dbb60306014afb80c2f03ba367ef9c9dd16b09ae Mon Sep 17 00:00:00 2001 From: Michael Chin Date: Thu, 16 Oct 2025 20:56:39 -0700 Subject: [PATCH 2/3] remove debug --- libs/aws/langchain_aws/chat_models/bedrock_converse.py | 6 ------ .../integration_tests/chat_models/test_bedrock_converse.py | 3 --- 2 files changed, 9 deletions(-) diff --git a/libs/aws/langchain_aws/chat_models/bedrock_converse.py b/libs/aws/langchain_aws/chat_models/bedrock_converse.py index 7333d7a5..3167993f 100644 --- a/libs/aws/langchain_aws/chat_models/bedrock_converse.py +++ b/libs/aws/langchain_aws/chat_models/bedrock_converse.py @@ -972,7 +972,6 @@ def _get_llm_for_structured_output_no_tool_choice( self, schema: Union[Dict, type], ) -> Runnable[LanguageModelInput, BaseMessage]: - print("In _get_llm_for_structured_output_no_tool_choice") admonition = ( "ChatBedrockConverse structured output relies on forced tool calling, " "which is not supported for this model. This method will raise " @@ -991,10 +990,8 @@ def _get_llm_for_structured_output_no_tool_choice( "by disabling `thinking`." ) admonition = f"{admonition} {additional_context}" - print("Returning warning for thinking claude model") warnings.warn(admonition) try: - print("Attempting bind_tools with structured output") llm = self.bind_tools( [schema], ls_structured_output_format={ @@ -1003,14 +1000,11 @@ def _get_llm_for_structured_output_no_tool_choice( }, ) except Exception: - print("Failed, using regular bind_tools") llm = self.bind_tools([schema]) def _raise_if_no_tool_calls(message: AIMessage) -> AIMessage: if not message.tool_calls: - print("No tool calls found, raising OutputParserException") raise OutputParserException(admonition) - print("Tool calls found, returning normally") return message return llm | _raise_if_no_tool_calls diff --git a/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py b/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py index c45355a3..4c05cbd7 100644 --- a/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py +++ b/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py @@ -454,14 +454,11 @@ def test_structured_output_tool_choice_not_supported(thinking_model: str) -> Non "thinking": {"type": "enabled", "budget_tokens": 2000} }, ) - print("Generating structured llm with thinking...") with pytest.warns(match="structured output"): structured_llm = llm.with_structured_output(ClassifyQuery) - print("Invoking on structured llm with thinking and related query...") response = structured_llm.invoke("How big are cats?") assert isinstance(response, ClassifyQuery) - print("Invoking on structured llm with thinking and generic query...") with pytest.raises(OutputParserException): structured_llm.invoke("Hello!") From a3158d13b063feed12571718a1e54cc9327f4f56 Mon Sep 17 00:00:00 2001 From: Michael Chin Date: Thu, 16 Oct 2025 20:58:25 -0700 Subject: [PATCH 3/3] lint --- .../integration_tests/chat_models/test_bedrock_converse.py | 2 +- libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py b/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py index 4c05cbd7..7cc6b416 100644 --- a/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py +++ b/libs/aws/tests/integration_tests/chat_models/test_bedrock_converse.py @@ -434,7 +434,7 @@ def test_guardrails() -> None: "us.anthropic.claude-opus-4-1-20250805-v1:0", "us.anthropic.claude-sonnet-4-5-20250929-v1:0", # - #"us.anthropic.claude-haiku-4-5-20251001-v1:0", + # "us.anthropic.claude-haiku-4-5-20251001-v1:0", ], ) def test_structured_output_tool_choice_not_supported(thinking_model: str) -> None: diff --git a/libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py b/libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py index b78c783a..bc678195 100644 --- a/libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py +++ b/libs/aws/tests/unit_tests/chat_models/test_bedrock_converse.py @@ -120,7 +120,7 @@ def test_anthropic_bind_tools_tool_choice() -> None: "anthropic.claude-3-7-sonnet-20250219-v1:0", "anthropic.claude-sonnet-4-20250514-v1:0", "anthropic.claude-opus-4-20250514-v1:0", - "anthropic.claude-haiku-4-5-20251001-v1:0" + "anthropic.claude-haiku-4-5-20251001-v1:0", ], ) def test_anthropic_thinking_bind_tools_tool_choice(thinking_model: str) -> None: