diff --git a/libs/aws/langchain_aws/chat_models/bedrock_converse.py b/libs/aws/langchain_aws/chat_models/bedrock_converse.py index fd40594f..cefd7610 100644 --- a/libs/aws/langchain_aws/chat_models/bedrock_converse.py +++ b/libs/aws/langchain_aws/chat_models/bedrock_converse.py @@ -546,6 +546,46 @@ class Joke(BaseModel): request_metadata: Optional[Dict[str, str]] = None """Key-Value pairs that you can use to filter invocation logs.""" + reasoning_effort: Optional[Literal["low", "medium", "high"]] = None + """Reasoning effort level for models that support extended thinking. + + Controls the computational effort used in the reasoning process. + Valid options are "low", "medium", or "high". + + This parameter provides a convenient way to enable reasoning without + manually configuring `additional_model_request_fields`. When set, it + will automatically configure the appropriate reasoning parameters for + the model provider: + + - **Amazon Nova models**: Configures `reasoningConfig` with the specified + `maxReasoningEffort` level. + - **OpenAI models on Bedrock**: Sets the `reasoning_effort` parameter. + + Example: + .. code-block:: python + + from langchain_aws import ChatBedrockConverse + + # Using reasoning_effort with Amazon Nova + llm = ChatBedrockConverse( + model="us.amazon.nova-lite-v1:0", + region_name="us-west-2", + reasoning_effort="medium", + ) + + # Using reasoning_effort with OpenAI on Bedrock + llm = ChatBedrockConverse( + model="openai.gpt-4o-mini-2024-07-18-v1:0", + region_name="us-west-2", + reasoning_effort="high", + ) + + Note: + For Anthropic Claude models, use `additional_model_request_fields` + with the `thinking` parameter instead, as Claude uses `budget_tokens` + rather than effort levels. + """ + guard_last_turn_only: bool = False """Boolean flag for applying the guardrail to only the last turn.""" @@ -861,6 +901,50 @@ def _set_model_profile(self) -> Self: self.profile = _get_default_model_profile(model_id) return self + @model_validator(mode="after") + def _configure_reasoning_effort(self) -> Self: + """Configure reasoning parameters based on reasoning_effort setting.""" + if self.reasoning_effort is None: + return self + + # Build the appropriate reasoning config based on provider + reasoning_config: Dict[str, Any] = {} + + if self.provider == "amazon": + # Amazon Nova models use reasoningConfig + reasoning_config = { + "reasoningConfig": { + "type": "enabled", + "maxReasoningEffort": self.reasoning_effort, + } + } + elif self.provider == "openai": + # OpenAI models on Bedrock use reasoning_effort + reasoning_config = {"reasoning_effort": self.reasoning_effort} + else: + # For other providers, warn that reasoning_effort may not be supported + warnings.warn( + f"reasoning_effort parameter may not be supported for provider " + f"'{self.provider}'. For Anthropic Claude models, use " + f"additional_model_request_fields with the 'thinking' parameter " + f"instead. The reasoning_effort value will be passed as-is to " + f"additionalModelRequestFields.", + UserWarning, + stacklevel=2, + ) + reasoning_config = {"reasoning_effort": self.reasoning_effort} + + # Merge with existing additional_model_request_fields + if self.additional_model_request_fields: + self.additional_model_request_fields = { + **reasoning_config, + **self.additional_model_request_fields, + } + else: + self.additional_model_request_fields = reasoning_config + + return self + def _get_base_model(self) -> str: """Return base model id, stripping any regional prefix.""" 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 d6211444..59d84143 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 @@ -2356,3 +2356,107 @@ def test_get_num_tokens_from_messages_api_error_fallback() -> None: token_count = llm.get_num_tokens_from_messages(messages) assert token_count == 5 mock_base.assert_called_once() + + +def test_reasoning_effort_amazon_nova() -> None: + """Test that reasoning_effort configures reasoningConfig for Amazon Nova models.""" + llm = ChatBedrockConverse( + model="us.amazon.nova-lite-v1:0", + region_name="us-west-2", + reasoning_effort="medium", + ) + + assert llm.additional_model_request_fields is not None + assert "reasoningConfig" in llm.additional_model_request_fields + assert llm.additional_model_request_fields["reasoningConfig"] == { + "type": "enabled", + "maxReasoningEffort": "medium", + } + + +@pytest.mark.parametrize("effort", ["low", "medium", "high"]) +def test_reasoning_effort_amazon_nova_all_levels( + effort: Literal["low", "medium", "high"] +) -> None: + """Test all reasoning effort levels for Amazon Nova models.""" + llm = ChatBedrockConverse( + model="amazon.nova-pro-v1:0", + region_name="us-west-2", + reasoning_effort=effort, + ) + + assert llm.additional_model_request_fields is not None + assert llm.additional_model_request_fields["reasoningConfig"]["maxReasoningEffort"] == effort + + +def test_reasoning_effort_openai() -> None: + """Test that reasoning_effort is set directly for OpenAI models on Bedrock.""" + llm = ChatBedrockConverse( + model="openai.gpt-oss-120b-1:0", + region_name="us-west-2", + reasoning_effort="high", + ) + + assert llm.additional_model_request_fields is not None + assert llm.additional_model_request_fields.get("reasoning_effort") == "high" + # Should not have reasoningConfig for OpenAI + assert "reasoningConfig" not in llm.additional_model_request_fields + + +def test_reasoning_effort_preserves_existing_fields() -> None: + """Test that reasoning_effort merges with existing additional_model_request_fields.""" + llm = ChatBedrockConverse( + model="us.amazon.nova-lite-v1:0", + region_name="us-west-2", + reasoning_effort="low", + additional_model_request_fields={"customParam": "value"}, + ) + + assert llm.additional_model_request_fields is not None + # Both the reasoning config and custom param should be present + assert "reasoningConfig" in llm.additional_model_request_fields + assert llm.additional_model_request_fields["customParam"] == "value" + + +def test_reasoning_effort_existing_fields_take_precedence() -> None: + """Test that existing additional_model_request_fields take precedence.""" + llm = ChatBedrockConverse( + model="us.amazon.nova-lite-v1:0", + region_name="us-west-2", + reasoning_effort="low", + additional_model_request_fields={ + "reasoningConfig": {"type": "disabled"}, + }, + ) + + assert llm.additional_model_request_fields is not None + # User-provided config should take precedence + assert llm.additional_model_request_fields["reasoningConfig"] == {"type": "disabled"} + + +def test_reasoning_effort_none_does_not_modify_fields() -> None: + """Test that reasoning_effort=None doesn't add any fields.""" + llm = ChatBedrockConverse( + model="us.amazon.nova-lite-v1:0", + region_name="us-west-2", + reasoning_effort=None, + ) + + assert llm.additional_model_request_fields is None + + +def test_reasoning_effort_unsupported_provider_warning() -> None: + """Test that using reasoning_effort with unsupported providers emits a warning.""" + with pytest.warns( + UserWarning, + match="reasoning_effort parameter may not be supported for provider", + ): + llm = ChatBedrockConverse( + model="anthropic.claude-3-sonnet-20240229-v1:0", + region_name="us-west-2", + reasoning_effort="medium", + ) + + # The value should still be passed through + assert llm.additional_model_request_fields is not None + assert llm.additional_model_request_fields.get("reasoning_effort") == "medium"