fix: applies Responses API validation to Strands, renames Strands tests

Signed-off-by: RanjitR <[email protected]>

Author: ronjer30

examples/frameworks/strands_demo/src/nat_strands_demo/configs/eval_config.yml

@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+# path-check-skip-file
 
 
 functions:

examples/frameworks/strands_demo/src/nat_strands_demo/configs/sizing_config.yml

@@ -29,6 +29,7 @@
 from nat.llm.utils.thinking import FunctionArgumentWrapper
 from nat.llm.utils.thinking import patch_with_thinking
 from nat.utils.exception_handlers.automatic_retries import patch_with_retry
+from nat.utils.responses_api import validate_no_responses_api
 from nat.utils.type_utils import override
 
 ModelType = TypeVar("ModelType")
@@ -95,9 +96,11 @@ def inject(self, messages, *args, **kwargs) -> FunctionArgumentWrapper:
 @register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.STRANDS)
 async def openai_strands(llm_config: OpenAIModelConfig, _builder: Builder):
 
+    validate_no_responses_api(llm_config, LLMFrameworkEnum.STRANDS)
+
     from strands.models.openai import OpenAIModel
 
-    params = llm_config.model_dump(exclude={"type", "api_key", "base_url", "model_name"},
+    params = llm_config.model_dump(exclude={"type", "api_type", "api_key", "base_url", "model_name"},
                                    by_alias=True,
                                    exclude_none=True)
     # Remove NAT-specific and retry-specific keys not accepted by OpenAI chat.create
@@ -119,6 +122,8 @@ async def openai_strands(llm_config: OpenAIModelConfig, _builder: Builder):
 @register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.STRANDS)
 async def nim_strands(llm_config: NIMModelConfig, _builder: Builder):
 
+    validate_no_responses_api(llm_config, LLMFrameworkEnum.STRANDS)
+
     # NIM is OpenAI compatible; use OpenAI model with NIM base_url and api_key
     from strands.models.openai import OpenAIModel
 
@@ -151,7 +156,7 @@ def format_request_messages(cls, messages, system_prompt=None):
 
             return formatted_messages
 
-    params = llm_config.model_dump(exclude={"type", "api_key", "base_url", "model_name"},
+    params = llm_config.model_dump(exclude={"type", "api_type", "api_key", "base_url", "model_name"},
                                    by_alias=True,
                                    exclude_none=True)
     # Remove NAT-specific and retry-specific keys not accepted by OpenAI
@@ -173,15 +178,26 @@ def format_request_messages(cls, messages, system_prompt=None):
 @register_llm_client(config_type=AWSBedrockModelConfig, wrapper_type=LLMFrameworkEnum.STRANDS)
 async def bedrock_strands(llm_config: AWSBedrockModelConfig, _builder: Builder):
 
+    validate_no_responses_api(llm_config, LLMFrameworkEnum.STRANDS)
+
     from strands.models.bedrock import BedrockModel
 
-    client = BedrockModel(
-        model_id=llm_config.model_name,
-        max_tokens=llm_config.max_tokens,
-        temperature=llm_config.temperature,
-        top_p=llm_config.top_p,
-        region_name=llm_config.region_name,
-        endpoint_url=llm_config.base_url,
-    )
+    params = llm_config.model_dump(exclude={"type", "api_type", "model_name", "region_name", "base_url"},
+                                   by_alias=True,
+                                   exclude_none=True)
+
+    for k in ("max_retries",
+              "num_retries",
+              "retry_on_status_codes",
+              "retry_on_errors",
+              "thinking",
+              "context_size",
+              "credentials_profile_name"):
+        params.pop(k, None)
+
+    client = BedrockModel(model_id=llm_config.model_name,
+                          region_name=llm_config.region_name,
+                          endpoint_url=llm_config.base_url,
+                          **params)
 
     yield _patch_llm_based_on_config(client, llm_config)

examples/frameworks/strands_demo/src/nat_strands_demo/configs/tracing_config.yml

@@ -276,13 +276,12 @@ def wrapped_init(agent_self, *args, **kwargs):
                 try:
                     # Import hook event types
                     # pylint: disable=import-outside-toplevel
-                    from strands.experimental.hooks import AfterToolInvocationEvent
-                    from strands.experimental.hooks import BeforeToolInvocationEvent
+                    from strands.hooks import AfterToolCallEvent
+                    from strands.hooks import BeforeToolCallEvent
 
                     # Register tool hooks on this agent instance
-                    agent_self.hooks.add_callback(BeforeToolInvocationEvent,
-                                                  handler.tool_hook.on_before_tool_invocation)
-                    agent_self.hooks.add_callback(AfterToolInvocationEvent, handler.tool_hook.on_after_tool_invocation)
+                    agent_self.hooks.add_callback(BeforeToolCallEvent, handler.tool_hook.on_before_tool_invocation)
+                    agent_self.hooks.add_callback(AfterToolCallEvent, handler.tool_hook.on_after_tool_invocation)
 
                     logger.debug("Strands tool hooks registered on Agent instance")
 

packages/nvidia_nat_strands/src/nat/plugins/strands/llm.py

@@ -582,20 +582,20 @@ def __init__(self, *args, **kwargs):
         def import_side_effect(module_name):
             if "agent.agent" in module_name:
                 return mock_agent_mod
-            elif "experimental.hooks" in module_name:
+            elif "hooks" in module_name and "strands" in module_name:
                 # Import the actual hook classes for testing
                 try:
-                    from strands.experimental.hooks import AfterToolInvocationEvent
-                    from strands.experimental.hooks import BeforeToolInvocationEvent
+                    from strands.hooks import AfterToolCallEvent
+                    from strands.hooks import BeforeToolCallEvent
                     hook_mod = MagicMock()
-                    hook_mod.BeforeToolInvocationEvent = BeforeToolInvocationEvent
-                    hook_mod.AfterToolInvocationEvent = AfterToolInvocationEvent
+                    hook_mod.BeforeToolCallEvent = BeforeToolCallEvent
+                    hook_mod.AfterToolCallEvent = AfterToolCallEvent
                     return hook_mod
                 except ImportError:
                     # Fallback to mocks if strands not available
                     hook_mod = MagicMock()
-                    hook_mod.BeforeToolInvocationEvent = MagicMock()
-                    hook_mod.AfterToolInvocationEvent = MagicMock()
+                    hook_mod.BeforeToolCallEvent = MagicMock()
+                    hook_mod.AfterToolCallEvent = MagicMock()
                     return hook_mod
             raise ImportError(f"No module named {module_name}")
 
@@ -663,10 +663,10 @@ def __init__(self, *args, **kwargs):
         def import_side_effect(module_name):
             if "agent.agent" in module_name:
                 return mock_agent_mod
-            elif "experimental.hooks" in module_name:
+            elif "hooks" in module_name and "strands" in module_name:
                 hook_mod = MagicMock()
-                hook_mod.BeforeToolInvocationEvent = MagicMock()
-                hook_mod.AfterToolInvocationEvent = MagicMock()
+                hook_mod.BeforeToolCallEvent = MagicMock()
+                hook_mod.AfterToolCallEvent = MagicMock()
                 return hook_mod
             raise ImportError(f"No module named {module_name}")
 

packages/nvidia_nat_strands/src/nat/plugins/strands/strands_callback_handler.py

@@ -19,6 +19,7 @@
 import pytest
 
 from nat.builder.builder import Builder
+from nat.data_models.llm import APITypeEnum
 from nat.llm.aws_bedrock_llm import AWSBedrockModelConfig
 from nat.llm.nim_llm import NIMModelConfig
 from nat.llm.openai_llm import OpenAIModelConfig
@@ -41,6 +42,11 @@ def openai_config(self):
         """Create an OpenAIModelConfig instance."""
         return OpenAIModelConfig(model_name="gpt-4")
 
+    @pytest.fixture
+    def openai_config_wrong_api(self):
+        """Create an OpenAIModelConfig with wrong API type."""
+        return OpenAIModelConfig(model_name="gpt-4", api_type=APITypeEnum.RESPONSES)
+
     @patch("strands.models.openai.OpenAIModel")
     async def test_openai_strands_basic(self, mock_model, openai_config, mock_builder):
         """Test that openai_strands as async context manager."""
@@ -64,6 +70,14 @@ async def test_openai_strands_with_params(self, mock_model, openai_config, mock_
         async with openai_strands(openai_config, mock_builder):
             mock_model.assert_called_once()
 
+    @patch("strands.models.openai.OpenAIModel")
+    async def test_api_type_validation(self, mock_model, openai_config_wrong_api, mock_builder):
+        """Non-chat-completion API types must raise a ValueError."""
+        with pytest.raises(ValueError):
+            async with openai_strands(openai_config_wrong_api, mock_builder):
+                pass
+        mock_model.assert_not_called()
+
 
 class TestBedrockStrands:
     """Tests for the bedrock_strands function."""
@@ -81,6 +95,15 @@ def bedrock_config(self):
             region_name="us-east-1",
         )
 
+    @pytest.fixture
+    def bedrock_config_wrong_api(self):
+        """Create an AWSBedrockModelConfig with wrong API type."""
+        return AWSBedrockModelConfig(
+            model_name="anthropic.claude-3-sonnet-20240229-v1:0",
+            region_name="us-east-1",
+            api_type=APITypeEnum.RESPONSES,
+        )
+
     @patch("strands.models.bedrock.BedrockModel")
     async def test_bedrock_strands_basic(self, mock_model, bedrock_config, mock_builder):
         """Test that bedrock_strands creates a BedrockModel."""
@@ -91,6 +114,14 @@ async def test_bedrock_strands_basic(self, mock_model, bedrock_config, mock_buil
         async with bedrock_strands(bedrock_config, mock_builder):
             mock_model.assert_called_once()
 
+    @patch("strands.models.bedrock.BedrockModel")
+    async def test_api_type_validation(self, mock_model, bedrock_config_wrong_api, mock_builder):
+        """Non-chat-completion API types must raise a ValueError."""
+        with pytest.raises(ValueError):
+            async with bedrock_strands(bedrock_config_wrong_api, mock_builder):
+                pass
+        mock_model.assert_not_called()
+
 
 class TestNIMStrands:
     """Tests for the nim_strands function."""
@@ -109,6 +140,16 @@ def nim_config(self):
             base_url="https://integrate.api.nvidia.com/v1",
         )
 
+    @pytest.fixture
+    def nim_config_wrong_api(self):
+        """Create a NIMModelConfig with wrong API type."""
+        return NIMModelConfig(
+            model_name="meta/llama-3.1-8b-instruct",
+            api_key="test-api-key",
+            base_url="https://integrate.api.nvidia.com/v1",
+            api_type=APITypeEnum.RESPONSES,
+        )
+
     async def test_nim_strands_basic(self, nim_config, mock_builder):
         """Test that nim_strands creates a NIMCompatibleOpenAIModel."""
         # Patch OpenAIModel.__init__ to track the call
@@ -220,6 +261,14 @@ async def test_nim_strands_excludes_nat_specific_params(self, mock_builder):
                 assert "thinking" not in params
                 assert "retry_on_status_codes" not in params
 
+    async def test_api_type_validation(self, nim_config_wrong_api, mock_builder):
+        """Non-chat-completion API types must raise a ValueError."""
+        with patch("strands.models.openai.OpenAIModel.__init__", return_value=None) as mock_init:
+            with pytest.raises(ValueError):
+                async with nim_strands(nim_config_wrong_api, mock_builder):
+                    pass
+            mock_init.assert_not_called()
+
 
 class TestPatchLLMBasedOnConfig:
     """Tests for _patch_llm_based_on_config function."""