nvidia-nat-strands 1.4.0a20251209__py3-none-any.whl

nat/meta/pypi.md

@@ -0,0 +1,23 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+-->
+
+![NVIDIA NeMo Agent Toolkit](https://media.githubusercontent.com/media/NVIDIA/NeMo-Agent-Toolkit/refs/heads/main/docs/source/_static/banner.png "NeMo Agent toolkit banner image")
+
+# NVIDIA NeMo Agent Toolkit Subpackage
+This is a subpackage for AWS Strands integration in NeMo Agent toolkit.
+
+For more information about the NVIDIA NeMo Agent toolkit, please visit the [NeMo Agent toolkit GitHub Repo](https://github.com/NVIDIA/NeMo-Agent-Toolkit).

nat/plugins/strands/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.

nat/plugins/strands/llm.py

@@ -0,0 +1,339 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""LLM provider wrappers for AWS Strands integration with NVIDIA NeMo Agent toolkit.
+
+This module provides Strands-compatible LLM client wrappers for the following providers:
+
+Supported Providers
+-------------------
+- **OpenAI**: Direct OpenAI API integration through ``OpenAIModelConfig``
+- **NVIDIA NIM**: OpenAI-compatible endpoints for NVIDIA models through ``NIMModelConfig``
+- **AWS Bedrock**: Amazon Bedrock models (such as Claude) through ``AWSBedrockModelConfig``
+
+Each wrapper:
+- Validates that Responses API features are disabled (Strands manages tool execution)
+- Patches clients with NeMo Agent toolkit retry logic from ``RetryMixin``
+- Injects chain-of-thought prompts when ``ThinkingMixin`` is configured
+- Removes NeMo Agent toolkit-specific config keys before instantiating Strands clients
+
+Future Provider Support
+-----------------------
+The following providers are not yet supported but could be contributed:
+
+- **Azure OpenAI**: Would require a Strands Azure OpenAI client wrapper similar to the
+  existing OpenAI integration. Contributors should follow the pattern established in
+  ``openai_strands`` and ensure Azure-specific authentication (endpoint, API version,
+  deployment name) is properly handled.
+
+- **LiteLLM**: The wrapper  would need to handle LiteLLM's unified interface across
+  multiple providers while preserving Strands' tool execution semantics.
+
+See the Strands documentation at https://strandsagents.com for model provider details.
+"""
+
+import os
+from collections.abc import AsyncGenerator
+from typing import Any
+from typing import TypeVar
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_llm_client
+from nat.data_models.common import get_secret_value
+from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.retry_mixin import RetryMixin
+from nat.data_models.thinking_mixin import ThinkingMixin
+from nat.llm.aws_bedrock_llm import AWSBedrockModelConfig
+from nat.llm.nim_llm import NIMModelConfig
+from nat.llm.openai_llm import OpenAIModelConfig
+from nat.llm.utils.thinking import BaseThinkingInjector
+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")
+
+
+def _patch_llm_based_on_config(client: ModelType, llm_config: LLMBaseConfig) -> ModelType:
+    """Patch a Strands client per NAT config (retries/thinking) and return it.
+
+    Args:
+        client: Concrete Strands model client instance.
+        llm_config: NAT LLM config with Retry/Thinking mixins.
+
+    Returns:
+        The patched client instance.
+    """
+
+    class StrandsThinkingInjector(BaseThinkingInjector):
+
+        @override
+        def inject(self, messages, *args, **kwargs) -> FunctionArgumentWrapper:
+            thinking_prompt = self.system_prompt
+            if not thinking_prompt:
+                return FunctionArgumentWrapper(messages, *args, **kwargs)
+
+            # Strands calls: model.stream(messages, tool_specs, system_prompt)
+            # So system_prompt is the 3rd positional argument (index 1 in *args)
+            new_args = list(args)
+            new_kwargs = dict(kwargs)
+
+            # Check if system_prompt is passed as positional argument
+            if len(new_args) >= 2:  # tool_specs, system_prompt
+                existing_system_prompt = new_args[1] or ""  # system_prompt
+                if existing_system_prompt:
+                    # Prepend thinking prompt to existing system prompt
+                    combined_prompt = f"{thinking_prompt}\n\n{existing_system_prompt}"
+                else:
+                    combined_prompt = thinking_prompt
+                new_args[1] = combined_prompt
+            elif "system_prompt" in new_kwargs:
+                # system_prompt passed as keyword argument
+                existing_system_prompt = new_kwargs["system_prompt"] or ""
+                if existing_system_prompt:
+                    combined_prompt = f"{thinking_prompt}\n\n{existing_system_prompt}"
+                else:
+                    combined_prompt = thinking_prompt
+                new_kwargs["system_prompt"] = combined_prompt
+            else:
+                # No system_prompt provided, add as keyword argument
+                new_kwargs["system_prompt"] = thinking_prompt
+
+            return FunctionArgumentWrapper(messages, *new_args, **new_kwargs)
+
+    if isinstance(llm_config, RetryMixin):
+        client = patch_with_retry(client,
+                                  retries=llm_config.num_retries,
+                                  retry_codes=llm_config.retry_on_status_codes,
+                                  retry_on_messages=llm_config.retry_on_errors)
+
+    if isinstance(llm_config, ThinkingMixin) and llm_config.thinking_system_prompt is not None:
+        client = patch_with_thinking(
+            client,
+            StrandsThinkingInjector(
+                system_prompt=llm_config.thinking_system_prompt,
+                function_names=[
+                    "stream",
+                    "structured_output",
+                ],
+            ))
+
+    return client
+
+
+@register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.STRANDS)
+async def openai_strands(llm_config: OpenAIModelConfig, _builder: Builder) -> AsyncGenerator[Any, None]:
+    """Build a Strands OpenAI client from an NVIDIA NeMo Agent toolkit configuration.
+
+    The wrapper requires the ``nvidia-nat[strands]`` extra and a valid OpenAI-compatible
+    API key. When ``llm_config.api_key`` is empty, the integration falls back to the
+    ``OPENAI_API_KEY`` environment variable. Responses API features are disabled through
+    ``validate_no_responses_api`` because Strands handles tool execution inside the
+    framework runtime. The yielded client is patched with NeMo Agent toolkit retry and
+    thinking hooks so that framework-level policies remain consistent.
+
+    Args:
+        llm_config: OpenAI configuration declared in the workflow.
+        _builder: Builder instance provided by the workflow factory (unused).
+
+    Yields:
+        Strands ``OpenAIModel`` objects ready to stream responses with NeMo Agent toolkit
+        retry/thinking behaviors applied.
+    """
+
+    validate_no_responses_api(llm_config, LLMFrameworkEnum.STRANDS)
+
+    from strands.models.openai import OpenAIModel
+
+    params = llm_config.model_dump(
+        exclude={"type", "api_type", "api_key", "base_url", "model_name", "max_retries", "thinking"},
+        by_alias=True,
+        exclude_none=True)
+
+    client = OpenAIModel(
+        client_args={
+            "api_key": get_secret_value(llm_config.api_key) or os.getenv("OPENAI_API_KEY"),
+            "base_url": llm_config.base_url,
+        },
+        model_id=llm_config.model_name,
+        params=params,
+    )
+
+    yield _patch_llm_based_on_config(client, llm_config)
+
+
+@register_llm_client(config_type=NIMModelConfig, wrapper_type=LLMFrameworkEnum.STRANDS)
+async def nim_strands(llm_config: NIMModelConfig, _builder: Builder) -> AsyncGenerator[Any, None]:
+    """Build a Strands OpenAI-compatible client for NVIDIA NIM endpoints.
+
+    Install the ``nvidia-nat[strands]`` extra and provide a NIM API key either through
+    ``llm_config.api_key`` or the ``NVIDIA_API_KEY`` environment variable. The wrapper
+    uses the OpenAI-compatible Strands client so Strands can route tool calls while the
+    NeMo Agent toolkit continues to manage retries, timeouts, and optional thinking
+    prompts. Responses API options are blocked to avoid conflicting execution models.
+
+    Args:
+        llm_config: Configuration for calling NVIDIA NIM by way of the OpenAI protocol.
+        _builder: Builder instance supplied during workflow construction (unused).
+
+    Yields:
+        Patched Strands clients that stream responses using the NVIDIA NIM endpoint
+        configured in ``llm_config``.
+    """
+
+    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
+
+    # Create a custom OpenAI model that formats text content as strings for NIM compatibility
+    class NIMCompatibleOpenAIModel(OpenAIModel):
+
+        @classmethod
+        def format_request_message_content(cls, content):
+            """Format OpenAI compatible content block with reasoning support.
+
+            Args:
+                content: Message content.
+
+            Returns:
+                OpenAI compatible content block.
+
+            Raises:
+                TypeError: If the content block type cannot be converted to
+                    an OpenAI-compatible format.
+            """
+            # Handle reasoning content by extracting the text
+            if isinstance(content, dict) and "reasoningContent" in content:
+                reasoning_text = content["reasoningContent"].get("reasoningText", {}).get("text", "")
+                return {"text": reasoning_text, "type": "text"}
+
+            # Fall back to parent implementation for other content types
+            return super().format_request_message_content(content)
+
+        @classmethod
+        def format_request_messages(cls, messages, system_prompt=None, *, system_prompt_content=None, **kwargs):
+            # Get the formatted messages from the parent
+            formatted_messages = super().format_request_messages(messages,
+                                                                 system_prompt,
+                                                                 system_prompt_content=system_prompt_content,
+                                                                 **kwargs)
+
+            # Convert content arrays with only text to strings for NIM
+            # compatibility
+            for msg in formatted_messages:
+                content = msg.get("content")
+                if (isinstance(content, list) and len(content) == 1 and isinstance(content[0], str)):
+                    # If content is a single-item list with a string, flatten it
+                    msg["content"] = content[0]
+                elif (isinstance(content, list)
+                      and all(isinstance(item, dict) and item.get("type") == "text" for item in content)):
+                    # If all items are text blocks, join them into a single
+                    # string
+                    text_content = "".join(item["text"] for item in content)
+                    # Ensure we don't send empty strings (NIM rejects them)
+                    msg["content"] = (text_content if text_content.strip() else " ")
+                elif isinstance(content, list) and len(content) == 0:
+                    # Handle empty content arrays
+                    msg["content"] = " "
+                elif isinstance(content, str) and not content.strip():
+                    # Handle empty strings
+                    msg["content"] = " "
+
+            return formatted_messages
+
+    params = llm_config.model_dump(
+        exclude={"type", "api_type", "api_key", "base_url", "model_name", "max_retries", "thinking"},
+        by_alias=True,
+        exclude_none=True)
+
+    # Determine base_url
+    base_url = llm_config.base_url or "https://integrate.api.nvidia.com/v1"
+
+    # Determine api_key; use dummy key for custom NIM endpoints without authentication
+    # If base_url is populated (not None) and no API key is available, use a dummy value
+    api_key = get_secret_value(llm_config.api_key) or os.getenv("NVIDIA_API_KEY")
+    if llm_config.base_url and llm_config.base_url.strip() and api_key is None:
+        api_key = "dummy-api-key"
+
+    client = NIMCompatibleOpenAIModel(
+        client_args={
+            "api_key": api_key,
+            "base_url": base_url,
+        },
+        model_id=llm_config.model_name,
+        params=params,
+    )
+
+    yield _patch_llm_based_on_config(client, llm_config)
+
+
+@register_llm_client(config_type=AWSBedrockModelConfig, wrapper_type=LLMFrameworkEnum.STRANDS)
+async def bedrock_strands(llm_config: AWSBedrockModelConfig, _builder: Builder) -> AsyncGenerator[Any, None]:
+    """Build a Strands Bedrock client from an NVIDIA NeMo Agent toolkit configuration.
+
+    The integration expects the ``nvidia-nat[strands]`` extra plus AWS credentials that
+    can be resolved by ``boto3``. Credentials are loaded in the following priority:
+
+    1. Explicit values embedded in the active AWS profile referenced by
+       ``llm_config.credentials_profile_name``.
+    2. Standard environment variables such as ``AWS_ACCESS_KEY_ID``,
+       ``AWS_SECRET_ACCESS_KEY``, and ``AWS_SESSION_TOKEN``.
+    3. Ambient credentials provided by the compute environment (for example, an IAM role
+       attached to the container or instance).
+
+    When ``llm_config.region_name`` is ``"None"`` or ``None`` Strands uses the regional
+    default configured in AWS. Responses API options remain unsupported so that Strands
+    can own tool execution. Retry and thinking hooks are added automatically before the
+    Bedrock client is yielded.
+
+    Args:
+        llm_config: AWS Bedrock configuration saved in the workflow.
+        _builder: Builder reference supplied by the workflow factory (unused).
+
+    Yields:
+        Strands ``BedrockModel`` instances configured for the selected Bedrock
+        ``model_name`` and patched with NeMo Agent toolkit retry/thinking helpers.
+    """
+
+    validate_no_responses_api(llm_config, LLMFrameworkEnum.STRANDS)
+
+    from strands.models.bedrock import BedrockModel
+
+    params = llm_config.model_dump(
+        exclude={
+            "type",
+            "api_type",
+            "model_name",
+            "region_name",
+            "base_url",
+            "max_retries",
+            "thinking",
+            "context_size",
+            "credentials_profile_name",
+        },
+        by_alias=True,
+        exclude_none=True,
+    )
+
+    region = None if llm_config.region_name in (None, "None") else llm_config.region_name
+    client = BedrockModel(model_id=llm_config.model_name,
+                          region_name=region,
+                          endpoint_url=llm_config.base_url,
+                          **params)
+
+    yield _patch_llm_based_on_config(client, llm_config)

nat/plugins/strands/register.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+# flake8: noqa
+# isort:skip_file
+
+from . import llm
+from . import tool_wrapper

nat/plugins/strands/strands_callback_handler.py

@@ -0,0 +1,615 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import asyncio
+import copy
+import importlib
+import json
+import logging
+import time
+import uuid
+from collections.abc import AsyncGenerator
+from collections.abc import Callable
+from typing import Any
+
+from nat.builder.context import Context
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.data_models.intermediate_step import IntermediateStepPayload
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.intermediate_step import StreamEventData
+from nat.data_models.intermediate_step import TraceMetadata
+from nat.data_models.intermediate_step import UsageInfo
+from nat.profiler.callbacks.base_callback_class import BaseProfilerCallback
+from nat.profiler.callbacks.token_usage_base_model import TokenUsageBaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class StrandsToolInstrumentationHook:
+    """Hook callbacks for instrumenting Strands tool invocations.
+
+    This class provides callbacks for Strands' hooks API to
+    capture tool execution events and emit proper TOOL_START/END spans.
+    """
+
+    def __init__(self, handler: 'StrandsProfilerHandler'):
+        """Initialize the hook with a reference to the profiler handler.
+
+        Args:
+            handler: StrandsProfilerHandler instance that manages this hook
+        """
+        self.handler = handler
+        self._tool_start_times: dict[str, float] = {}
+        self._step_manager = Context.get().intermediate_step_manager
+
+    def on_before_tool_invocation(self, event: Any) -> None:
+        """Handle tool invocation start.
+
+        Called by Strands before a tool is executed.
+        Emits a TOOL_START span.
+
+        Args:
+            event: BeforeToolInvocationEvent from Strands
+        """
+        try:
+            tool_use = event.tool_use
+            selected_tool = event.selected_tool
+
+            if not selected_tool:
+                logger.debug("Tool hook: no selected_tool, skipping")
+                return
+
+            # Extract tool information
+            tool_name, tool_use_id, tool_input = self._extract_tool_info(selected_tool, tool_use)
+
+            # Store start time for duration calculation
+            self._tool_start_times[tool_use_id] = time.time()
+
+            step_manager = self._step_manager
+
+            start_payload = IntermediateStepPayload(
+                event_type=IntermediateStepType.TOOL_START,
+                framework=LLMFrameworkEnum.STRANDS,
+                name=tool_name,
+                UUID=tool_use_id,
+                data=StreamEventData(input=str(tool_input), output=""),
+                metadata=TraceMetadata(
+                    tool_inputs=copy.deepcopy(tool_input),
+                    tool_info=copy.deepcopy(getattr(selected_tool, 'tool_spec', {})),
+                ),
+                usage_info=UsageInfo(token_usage=TokenUsageBaseModel()),
+            )
+
+            step_manager.push_intermediate_step(start_payload)
+
+            logger.debug("TOOL_START: %s (ID: %s)", tool_name, tool_use_id)
+        except Exception:  # noqa: BLE001
+            logger.error("Error in before_tool_invocation")
+            raise
+
+    def on_after_tool_invocation(self, event: Any) -> None:
+        """Handle tool invocation end.
+
+        Called by Strands after a tool execution completes.
+        Emits a TOOL_END span.
+
+        Args:
+            event: AfterToolInvocationEvent from Strands
+        """
+        try:
+            tool_use = event.tool_use
+            selected_tool = event.selected_tool
+            result = event.result
+            exception = event.exception
+
+            if not selected_tool:
+                logger.debug("Tool hook: no selected_tool, skipping")
+                return
+
+            # Extract tool information
+            tool_name, tool_use_id, tool_input = self._extract_tool_info(selected_tool, tool_use)
+            start_time = self._tool_start_times.pop(tool_use_id, time.time())
+
+            # Extract output from result
+            tool_output = ""
+            if isinstance(result, dict):
+                content = result.get('content', [])
+                if isinstance(content, list):
+                    for item in content:
+                        if isinstance(item, dict) and 'text' in item:
+                            tool_output += item['text']
+
+            # Handle errors
+            if exception:
+                tool_output = f"Error: {exception}"
+
+            # Use stored step_manager to avoid context isolation issues
+            step_manager = self._step_manager
+
+            end_payload = IntermediateStepPayload(
+                event_type=IntermediateStepType.TOOL_END,
+                span_event_timestamp=start_time,
+                framework=LLMFrameworkEnum.STRANDS,
+                name=tool_name,
+                UUID=tool_use_id,
+                metadata=TraceMetadata(tool_outputs=tool_output),
+                usage_info=UsageInfo(token_usage=TokenUsageBaseModel()),
+                data=StreamEventData(input=str(tool_input), output=tool_output),
+            )
+            step_manager.push_intermediate_step(end_payload)
+
+            logger.debug("TOOL_END: %s (ID: %s)", tool_name, tool_use_id)
+
+        except Exception:  # noqa: BLE001
+            logger.error("Failed to handle after_tool_invocation")
+            raise
+
+    def _extract_tool_info(self, selected_tool: Any, tool_use: dict) -> tuple[str, str, dict]:
+        """Extract tool name, ID, and input from event.
+
+        Args:
+            selected_tool: The tool being invoked
+            tool_use: Tool use dictionary from Strands event
+
+        Returns:
+            Tuple of (tool_name, tool_use_id, tool_input)
+        """
+        tool_name = getattr(selected_tool, 'tool_name', tool_use.get('name', 'unknown_tool'))
+        tool_use_id = tool_use.get('toolUseId')
+        if tool_use_id is None:
+            logger.warning("Missing toolUseId in tool_use event, using 'unknown' fallback")
+            tool_use_id = "unknown"
+        tool_input = tool_use.get('input', {}) or {}
+        return tool_name, tool_use_id, tool_input
+
+
+class StrandsProfilerHandler(BaseProfilerCallback):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._patched: bool = False
+        self.last_call_ts = time.time()
+
+        # Note: tool hooks are now created per-agent-instance in wrapped_init
+        # to avoid shared state in concurrent execution
+
+    def instrument(self) -> None:
+        """
+        Instrument Strands for telemetry capture.
+
+        This patches:
+        1. Model streaming methods (OpenAI/Bedrock) for LLM spans
+        2. Agent.__init__ to auto-register tool hooks on Agent creation
+
+        Tool instrumentation uses Strands' hooks API,
+        which is automatically registered when an Agent is instantiated.
+        """
+        if self._patched:
+            return
+
+        try:
+            # Patch LLM streaming methods
+            OpenAIModel = None
+            BedrockModel = None
+            try:
+                openai_mod = importlib.import_module("strands.models.openai")
+                OpenAIModel = getattr(openai_mod, "OpenAIModel", None)
+            except Exception:  # noqa: BLE001
+                OpenAIModel = None
+
+            try:
+                bedrock_mod = importlib.import_module("strands.models.bedrock")
+                BedrockModel = getattr(bedrock_mod, "BedrockModel", None)
+            except Exception:  # noqa: BLE001
+                BedrockModel = None
+
+            to_patch: list[tuple[type, str]] = []
+            if OpenAIModel is not None:
+                for name in ("stream", "structured_output"):
+                    if hasattr(OpenAIModel, name):
+                        to_patch.append((OpenAIModel, name))
+            if BedrockModel is not None:
+                for name in ("stream", "structured_output"):
+                    if hasattr(BedrockModel, name):
+                        to_patch.append((BedrockModel, name))
+
+            for cls, method_name in to_patch:
+                original = getattr(cls, method_name)
+                wrapped = self._wrap_stream_method(original)
+                setattr(cls, method_name, wrapped)
+
+            debug_targets = [f"{c.__name__}.{m}" for c, m in to_patch]
+            logger.info(
+                "StrandsProfilerHandler LLM instrumentation: %s",
+                debug_targets,
+            )
+
+            # Patch Agent.__init__ to auto-register hooks
+            self._instrument_agent_init()
+
+            self._patched = True
+
+        except Exception:  # noqa: BLE001
+            logger.error("Failed to instrument Strands models")
+            raise
+
+    def _instrument_agent_init(self) -> None:
+        """Patch Agent.__init__ to auto-register hooks on instantiation.
+
+        This ensures that whenever a Strands Agent is created, our tool
+        instrumentation hooks are automatically registered without requiring
+        any user code changes.
+        """
+        try:
+            # Import Agent class
+            agent_mod = importlib.import_module("strands.agent.agent")
+            Agent = getattr(agent_mod, "Agent", None)
+
+            if Agent is None:
+                logger.warning("Agent class not found in strands.agent.agent")
+                return
+
+            # Save reference to handler in closure
+            handler = self
+
+            # Save original __init__
+            original_init = Agent.__init__
+
+            def wrapped_init(agent_self, *args, **kwargs):
+                """Wrapped Agent.__init__ that auto-registers hooks."""
+                # Call original init
+                original_init(agent_self, *args, **kwargs)
+
+                # Auto-register tool hooks on this agent instance
+                try:
+                    # Import hook event types
+                    # pylint: disable=import-outside-toplevel
+                    from strands.hooks import AfterToolCallEvent
+                    from strands.hooks import BeforeToolCallEvent
+
+                    # Create a dedicated hook instance for this agent
+                    agent_tool_hook = StrandsToolInstrumentationHook(handler)
+
+                    # Register tool hooks on this agent instance
+                    agent_self.hooks.add_callback(BeforeToolCallEvent, agent_tool_hook.on_before_tool_invocation)
+                    agent_self.hooks.add_callback(AfterToolCallEvent, agent_tool_hook.on_after_tool_invocation)
+
+                    logger.debug("Strands tool hooks registered on Agent instance")
+
+                except Exception:  # noqa: BLE001
+                    logger.exception("Failed to auto-register hooks")
+
+            # Replace Agent.__init__ with wrapped version
+            Agent.__init__ = wrapped_init
+
+            logger.info("Strands Agent.__init__ instrumentation applied")
+
+        except Exception:  # noqa: BLE001
+            logger.exception("Failed to instrument Agent.__init__")
+
+    def _extract_model_info(self, model_instance: Any) -> tuple[str, dict[str, Any]]:
+        """Extract model name from Strands model instance."""
+        model_name = ""
+
+        for attr_name in ['config', 'client_args']:
+            if hasattr(model_instance, attr_name):
+                attr_value = getattr(model_instance, attr_name, None)
+                if isinstance(attr_value, dict):
+                    for key, val in attr_value.items():
+                        if 'model' in key.lower() and val:
+                            model_name = str(val)
+                            break
+                if model_name:
+                    break
+
+        return str(model_name), {}
+
+    def _wrap_stream_method(self, original: Callable[..., Any]) -> Callable[..., Any]:
+        # Capture handler reference in closure
+        handler = self
+
+        async def wrapped(model_self, *args, **kwargs) -> AsyncGenerator[Any, None]:  # type: ignore[override]
+            """
+            Wrapper for Strands model streaming that emits paired
+            LLM_START/END spans with usage and metrics.
+            """
+            context = Context.get()
+            step_manager = context.intermediate_step_manager
+
+            event_uuid = str(uuid.uuid4())
+            start_time = time.time()
+
+            # Extract model info and parameters
+            model_name, _ = handler._extract_model_info(model_self)
+
+            # Extract messages from args (Strands passes as positional args)
+            # Signature: stream(self, messages, tool_specs=None,
+            #                   system_prompt=None, **kwargs)
+            raw_messages = args[0] if args else []
+            tool_specs = args[1] if len(args) > 1 else kwargs.get("tool_specs")
+            system_prompt = (args[2] if len(args) > 2 else kwargs.get("system_prompt"))
+
+            # Build chat_inputs with system prompt and messages
+            all_messages = []
+            if system_prompt:
+                all_messages.append({"text": system_prompt, "role": "system"})
+            if isinstance(raw_messages, list):
+                all_messages.extend(copy.deepcopy(raw_messages))
+
+            # Extract tools schema for metadata
+            tools_schema = []
+            if tool_specs and isinstance(tool_specs, list):
+                try:
+                    tools_schema = [{
+                        "type": "function",
+                        "function": {
+                            "name": tool_spec.get("name", "unknown"),
+                            "description": tool_spec.get("description", ""),
+                            "parameters": tool_spec.get("inputSchema", {}).get("json", {})
+                        }
+                    } for tool_spec in tool_specs]
+                except Exception:  # noqa: BLE001
+                    logger.debug("Failed to extract tools schema", exc_info=True)
+                    tools_schema = []
+
+            # Extract string representation of last user message for data.input
+            # (full message history is in metadata.chat_inputs)
+            llm_input_str = ""
+            if all_messages:
+                last_msg = all_messages[-1]
+                if isinstance(last_msg, dict) and 'text' in last_msg:
+                    llm_input_str = last_msg['text']
+                elif isinstance(last_msg, dict):
+                    llm_input_str = str(last_msg)
+                else:
+                    llm_input_str = str(last_msg)
+
+            # Always emit START first (before streaming begins)
+            start_payload = IntermediateStepPayload(
+                event_type=IntermediateStepType.LLM_START,
+                framework=LLMFrameworkEnum.STRANDS,
+                name=str(model_name),
+                UUID=event_uuid,
+                data=StreamEventData(input=llm_input_str, output=""),
+                metadata=TraceMetadata(
+                    chat_inputs=copy.deepcopy(all_messages),
+                    tools_schema=copy.deepcopy(tools_schema),
+                ),
+                usage_info=UsageInfo(
+                    token_usage=TokenUsageBaseModel(),
+                    num_llm_calls=1,
+                    seconds_between_calls=int(time.time() - self.last_call_ts),
+                ),
+            )
+            step_manager.push_intermediate_step(start_payload)
+            self.last_call_ts = time.time()
+
+            # Collect output text, tool calls, and token usage while streaming
+            output_text = ""
+            tool_calls = []  # List of tool calls made by the LLM
+            current_tool_call = None  # Currently accumulating tool call
+            token_usage = TokenUsageBaseModel()
+            ended: bool = False
+
+            def _push_end_if_needed() -> None:
+                nonlocal ended
+                if ended:
+                    return
+
+                # Determine the output to show in the span
+                # If there are tool calls, format them as the output
+                # Otherwise, use the text response
+                if tool_calls:
+                    # Format tool calls as readable output
+                    tool_call_strs = []
+                    for tc in tool_calls:
+                        tool_name = tc.get('name', 'unknown')
+                        tool_input = tc.get('input', {})
+                        tool_call_strs.append(f"Tool: {tool_name}\nInput: {tool_input}")
+                    output_content = "\n\n".join(tool_call_strs)
+                else:
+                    output_content = output_text
+
+                chat_responses_list = []
+                if output_content:
+                    chat_responses_list = [output_content]
+
+                # Build metadata with standard NAT structure
+                metadata = TraceMetadata(
+                    chat_responses=chat_responses_list,
+                    chat_inputs=all_messages,
+                    tools_schema=copy.deepcopy(tools_schema),
+                )
+
+                # Push END with input/output and token usage
+                end_payload = IntermediateStepPayload(
+                    event_type=IntermediateStepType.LLM_END,
+                    span_event_timestamp=start_time,
+                    framework=LLMFrameworkEnum.STRANDS,
+                    name=str(model_name),
+                    UUID=event_uuid,
+                    data=StreamEventData(input=llm_input_str, output=output_content),
+                    usage_info=UsageInfo(token_usage=token_usage, num_llm_calls=1),
+                    metadata=metadata,
+                )
+                step_manager.push_intermediate_step(end_payload)
+                ended = True
+
+            try:
+                agen = original(model_self, *args, **kwargs)
+                if hasattr(agen, "__aiter__"):
+                    async for ev in agen:  # type: ignore
+                        try:
+                            # Extract text content
+                            text_content = self._extract_text_from_event(ev)
+                            if text_content:
+                                output_text += text_content
+
+                            # Extract tool call information
+                            tool_call_info = self._extract_tool_call_from_event(ev)
+                            if tool_call_info:
+                                if "name" in tool_call_info:
+                                    # New tool call starting
+                                    if current_tool_call:
+                                        # Finalize and save previous tool call
+                                        self._finalize_tool_call(current_tool_call)
+                                        tool_calls.append(current_tool_call)
+                                    current_tool_call = tool_call_info
+                                elif "input_chunk" in tool_call_info and current_tool_call:
+                                    # Accumulate input JSON string chunks
+                                    current_tool_call["input_str"] += tool_call_info["input_chunk"]
+
+                            # Check for contentBlockStop to finalize current tool call
+                            if "contentBlockStop" in ev and current_tool_call:
+                                self._finalize_tool_call(current_tool_call)
+                                tool_calls.append(current_tool_call)
+                                current_tool_call = None
+
+                            # Extract usage information (but don't push END yet - wait for all text)
+                            usage_info = self._extract_usage_from_event(ev)
+                            if usage_info:
+                                token_usage = TokenUsageBaseModel(**usage_info)
+
+                        except Exception:  # noqa: BLE001
+                            logger.debug("Failed to extract streaming fields from event", exc_info=True)
+                        yield ev
+                else:
+                    # Non-async generator fallback
+                    res = agen
+                    if asyncio.iscoroutine(res):
+                        res = await res  # type: ignore[func-returns-value]
+                    yield res
+            finally:
+                # Ensure END is always pushed
+                _push_end_if_needed()
+
+        return wrapped
+
+    def _extract_text_from_event(self, ev: dict) -> str:
+        """Extract text content from a Strands event.
+
+        Args:
+            ev: Event dictionary from Strands stream
+
+        Returns:
+            Extracted text content or empty string
+        """
+        if not isinstance(ev, dict):
+            return ""
+
+        # Try multiple possible locations for text content
+        if "data" in ev:
+            return str(ev["data"])
+
+        # Check for Strands contentBlockDelta structure (for streaming text responses)
+        if "contentBlockDelta" in ev and isinstance(ev["contentBlockDelta"], dict):
+            delta = ev["contentBlockDelta"].get("delta", {})
+            if isinstance(delta, dict) and "text" in delta:
+                return str(delta["text"])
+
+        # Check for other common text fields
+        if "content" in ev:
+            return str(ev["content"])
+
+        if "text" in ev:
+            return str(ev["text"])
+
+        # Check for nested content
+        if "message" in ev and isinstance(ev["message"], dict):
+            if "content" in ev["message"]:
+                return str(ev["message"]["content"])
+
+        return ""
+
+    def _finalize_tool_call(self, tool_call: dict[str, Any]) -> None:
+        """Parse the accumulated input_str JSON and store in the input field.
+
+        Args:
+            tool_call: Tool call dictionary with input_str to parse
+        """
+        input_str = tool_call.get("input_str", "")
+        if input_str:
+            try:
+                tool_call["input"] = json.loads(input_str)
+            except (json.JSONDecodeError, ValueError):
+                logger.debug("Failed to parse tool input JSON: %s", input_str)
+                tool_call["input"] = {"raw": input_str}
+        # Remove the temporary input_str field
+        tool_call.pop("input_str", None)
+
+    def _extract_tool_call_from_event(self, ev: dict) -> dict[str, Any] | None:
+        """Extract tool call information from a Strands event.
+
+        Args:
+            ev: Event dictionary from Strands stream
+
+        Returns:
+            Dictionary with tool call info (name, input_chunk) or None if not a tool call
+        """
+        if not isinstance(ev, dict):
+            return None
+
+        # Check for contentBlockStart with toolUse
+        if "contentBlockStart" in ev:
+            start = ev["contentBlockStart"].get("start", {})
+            if isinstance(start, dict) and "toolUse" in start:
+                tool_use = start["toolUse"]
+                return {
+                    "name": tool_use.get("name", "unknown"),
+                    "id": tool_use.get("toolUseId", "unknown"),
+                    "input_str": "",  # Will accumulate JSON string chunks
+                    "input": {}  # Will be parsed at the end
+                }
+
+        # Check for contentBlockDelta with toolUse input (streaming chunks)
+        if "contentBlockDelta" in ev:
+            delta = ev["contentBlockDelta"].get("delta", {})
+            if isinstance(delta, dict) and "toolUse" in delta:
+                tool_use_delta = delta["toolUse"]
+                input_chunk = tool_use_delta.get("input", "")
+                if input_chunk:
+                    # Return the chunk to be accumulated
+                    return {"input_chunk": input_chunk}
+
+        return None
+
+    def _extract_usage_from_event(self, ev: dict) -> dict[str, int] | None:
+        """Extract usage information from a Strands event.
+
+        Args:
+            ev: Event dictionary from Strands stream
+
+        Returns:
+            Dictionary with token usage info or None if not found
+        """
+        if not isinstance(ev, dict):
+            return None
+
+        md = ev.get("metadata")
+        if not isinstance(md, dict):
+            return None
+
+        usage = md.get("usage")
+        if not isinstance(usage, dict):
+            return None
+
+        try:
+            return {
+                "prompt_tokens": int(usage.get("inputTokens") or 0),
+                "completion_tokens": int(usage.get("outputTokens") or 0),
+                "total_tokens": int(usage.get("totalTokens") or 0),
+            }
+        except (ValueError, TypeError):
+            return None

nat/plugins/strands/tool_wrapper.py

@@ -0,0 +1,127 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import logging
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from pydantic import BaseModel
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function import Function
+from nat.cli.register_workflow import register_tool_wrapper
+from strands.types.tools import AgentTool  # type: ignore
+from strands.types.tools import ToolSpec  # type: ignore
+from strands.types.tools import ToolUse  # type: ignore
+
+logger = logging.getLogger(__name__)
+
+
+def _json_schema_from_pydantic(model: type[BaseModel]) -> dict[str, Any]:
+    try:
+        schema = model.model_json_schema()
+        for k in ("title", "additionalProperties"):
+            if k in schema:
+                del schema[k]
+        return {"json": schema}
+    except Exception:
+        logger.exception("Failed to generate JSON schema")
+        return {"json": {}}
+
+
+def _to_tool_result(tool_use_id: str, value: Any) -> dict[str, Any]:
+    if isinstance(value, (dict, list, tuple)):  # noqa: UP038
+        content_item = {"json": value}
+    else:
+        content_item = {"text": str(value)}
+    return {
+        "toolUseId": tool_use_id,
+        "status": "success",
+        "content": [content_item],
+    }
+
+
+def _to_error_result(tool_use_id: str, err: Exception) -> dict[str, Any]:
+    return {
+        "toolUseId": tool_use_id,
+        "status": "error",
+        "content": [{
+            "text": f"{type(err).__name__}: {err!s}"
+        }],
+    }
+
+
+class NATFunctionAgentTool(AgentTool):
+    """Concrete Strands AgentTool that wraps a NAT Function."""
+
+    def __init__(self, name: str, description: str | None, input_schema: dict[str, Any], fn: Function) -> None:
+        super().__init__()
+
+        self._tool_name = name
+        self._tool_spec: ToolSpec = {
+            "name": name,
+            "description": description or name,
+            "inputSchema": input_schema,
+        }
+        self._fn = fn
+
+    @property
+    def tool_name(self) -> str:
+        return self._tool_name
+
+    @property
+    def tool_spec(self) -> ToolSpec:
+        return self._tool_spec
+
+    @property
+    def tool_type(self) -> str:
+        return "function"
+
+    async def stream(self, tool_use: ToolUse, _invocation_state: dict[str, Any],
+                     **_kwargs: Any) -> AsyncGenerator[Any, None]:
+        from strands.types._events import ToolResultEvent  # type: ignore
+        from strands.types._events import ToolStreamEvent
+
+        tool_use_id = tool_use.get("toolUseId", "unknown")
+        tool_input = tool_use.get("input", {}) or {}
+
+        try:
+            if (self._fn.has_streaming_output and not self._fn.has_single_output):
+                last_chunk: Any | None = None
+                async for chunk in self._fn.acall_stream(**tool_input):
+                    last_chunk = chunk
+                    yield ToolStreamEvent(tool_use, chunk)
+                final = _to_tool_result(tool_use_id, last_chunk if last_chunk is not None else "")
+                yield ToolResultEvent(final)
+                return
+
+            result = await self._fn.acall_invoke(**tool_input)
+            yield ToolResultEvent(_to_tool_result(tool_use_id, result))
+        except Exception as exc:  # noqa: BLE001
+            logger.exception("Strands tool '%s' failed", self.tool_name)
+            yield ToolResultEvent(_to_error_result(tool_use_id, exc))
+
+
+@register_tool_wrapper(wrapper_type=LLMFrameworkEnum.STRANDS)
+def strands_tool_wrapper(name: str, fn: Function, _builder: Builder) -> NATFunctionAgentTool:
+    """Create a Strands `AgentTool` wrapper for a NAT `Function`."""
+    if fn.input_schema is None:
+        raise ValueError(f"Tool '{name}' must define an input schema")
+
+    input_schema = _json_schema_from_pydantic(fn.input_schema)
+    description = fn.description or name
+    return NATFunctionAgentTool(name=name, description=description, input_schema=input_schema, fn=fn)

nvidia_nat_strands-1.4.0a20251209.dist-info/METADATA

@@ -0,0 +1,43 @@
+Metadata-Version: 2.4
+Name: nvidia-nat-strands
+Version: 1.4.0a20251209
+Summary: Subpackage for AWS Strands integration in NeMo Agent toolkit
+Author: NVIDIA Corporation
+Maintainer: NVIDIA Corporation
+License: Apache-2.0
+Project-URL: documentation, https://docs.nvidia.com/nemo/agent-toolkit/latest/
+Project-URL: source, https://github.com/NVIDIA/NeMo-Agent-Toolkit
+Keywords: ai,rag,agents
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: <3.14,>=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: nvidia-nat==v1.4.0a20251209
+Requires-Dist: strands-agents~=1.17
+Requires-Dist: strands-agents-tools~=0.2
+
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+-->
+
+![NVIDIA NeMo Agent Toolkit](https://media.githubusercontent.com/media/NVIDIA/NeMo-Agent-Toolkit/refs/heads/main/docs/source/_static/banner.png "NeMo Agent toolkit banner image")
+
+# NVIDIA NeMo Agent Toolkit Subpackage
+This is a subpackage for AWS Strands integration in NeMo Agent toolkit.
+
+For more information about the NVIDIA NeMo Agent toolkit, please visit the [NeMo Agent toolkit GitHub Repo](https://github.com/NVIDIA/NeMo-Agent-Toolkit).

nvidia_nat_strands-1.4.0a20251209.dist-info/RECORD

@@ -0,0 +1,11 @@
+nat/meta/pypi.md,sha256=i_a-Zt6wbWAPjlFqa6CsvuaMZTjglgKxF7HzX3OZW5g,1112
+nat/plugins/strands/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
+nat/plugins/strands/llm.py,sha256=0vGhfS-98E0S2G6pPCRJ-Y_G53vCUofNvPwr-9mmL2U,15057
+nat/plugins/strands/register.py,sha256=oaVjGKwLkqzsSvimHHvgxfOLdpjEsEy8-VCtZPep5Bw,760
+nat/plugins/strands/strands_callback_handler.py,sha256=BeCsR4tTSxk7AVSt7Hq3KqEDyXI3Zwz2j4cwr_rAJv0,24992
+nat/plugins/strands/tool_wrapper.py,sha256=uWWEK4zMe3tL6j-Y_FQOAlD4rgw9O1V3iZxpt03JSk0,4604
+nvidia_nat_strands-1.4.0a20251209.dist-info/METADATA,sha256=OHR8MPzLKKyeG3AIUtBAWKZQyPSzfNByHJVF8Ky-4ZY,1887
+nvidia_nat_strands-1.4.0a20251209.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_strands-1.4.0a20251209.dist-info/entry_points.txt,sha256=6lerpn7DNmp8gJEPtjQIUA-fF6LlOhCm77lCP6ZOPA4,60
+nvidia_nat_strands-1.4.0a20251209.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_strands-1.4.0a20251209.dist-info/RECORD,,

nvidia_nat_strands-1.4.0a20251209.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

nvidia_nat_strands-1.4.0a20251209.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[nat.components]
+nat_strands = nat.plugins.strands.register

nvidia_nat_strands-1.4.0a20251209.dist-info/top_level.txt

@@ -0,0 +1 @@
+nat