nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.4.0a20251112__py3-none-any.whl

nat/front_ends/fastapi/main.py

@@ -13,19 +13,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import importlib
 import logging
 import os
+import typing
 
 from nat.front_ends.fastapi.fastapi_front_end_plugin_worker import FastApiFrontEndPluginWorkerBase
+from nat.front_ends.fastapi.utils import get_config_file_path
+from nat.front_ends.fastapi.utils import import_class_from_string
 from nat.runtime.loader import load_config
 
+if typing.TYPE_CHECKING:
+    from fastapi import FastAPI
+
 logger = logging.getLogger(__name__)
 
 
-def get_app():
+def get_app() -> "FastAPI":
 
-    config_file_path = os.getenv("NAT_CONFIG_FILE")
+    config_file_path = get_config_file_path()
     front_end_worker_full_name = os.getenv("NAT_FRONT_END_WORKER")
 
     if (not config_file_path):
@@ -36,28 +41,15 @@ def get_app():
 
     # Try to import the front end worker class
     try:
-        # Split the package from the class
-        front_end_worker_parts = front_end_worker_full_name.split(".")
-
-        front_end_worker_module_name = ".".join(front_end_worker_parts[:-1])
-        front_end_worker_class_name = front_end_worker_parts[-1]
-
-        front_end_worker_module = importlib.import_module(front_end_worker_module_name)
-
-        if not hasattr(front_end_worker_module, front_end_worker_class_name):
-            raise ValueError(f"Front end worker {front_end_worker_full_name} not found.")
-
-        front_end_worker_class: type[FastApiFrontEndPluginWorkerBase] = getattr(front_end_worker_module,
-                                                                                front_end_worker_class_name)
+        front_end_worker_class: type[FastApiFrontEndPluginWorkerBase] = import_class_from_string(
+            front_end_worker_full_name)
 
         if (not issubclass(front_end_worker_class, FastApiFrontEndPluginWorkerBase)):
             raise ValueError(
                 f"Front end worker {front_end_worker_full_name} is not a subclass of FastApiFrontEndPluginWorker.")
 
         # Load the config
-        abs_config_file_path = os.path.abspath(config_file_path)
-
-        config = load_config(abs_config_file_path)
+        config = load_config(config_file_path)
 
         # Create an instance of the front end worker class
         front_end_worker = front_end_worker_class(config)

nat/front_ends/fastapi/message_handler.py

@@ -25,6 +25,7 @@ from pydantic import ValidationError
 from starlette.websockets import WebSocketDisconnect
 
 from nat.authentication.interfaces import FlowHandlerBase
+from nat.data_models.api_server import ChatRequest
 from nat.data_models.api_server import ChatResponse
 from nat.data_models.api_server import ChatResponseChunk
 from nat.data_models.api_server import Error
@@ -33,6 +34,8 @@ from nat.data_models.api_server import ResponsePayloadOutput
 from nat.data_models.api_server import ResponseSerializable
 from nat.data_models.api_server import SystemResponseContent
 from nat.data_models.api_server import TextContent
+from nat.data_models.api_server import UserMessageContentRoleType
+from nat.data_models.api_server import UserMessages
 from nat.data_models.api_server import WebSocketMessageStatus
 from nat.data_models.api_server import WebSocketMessageType
 from nat.data_models.api_server import WebSocketSystemInteractionMessage
@@ -64,12 +67,12 @@ class WebSocketMessageHandler:
         self._running_workflow_task: asyncio.Task | None = None
         self._message_parent_id: str = "default_id"
         self._conversation_id: str | None = None
-        self._workflow_schema_type: str = None
-        self._user_interaction_response: asyncio.Future[HumanResponse] | None = None
+        self._workflow_schema_type: str | None = None
+        self._user_interaction_response: asyncio.Future[TextContent] | None = None
 
         self._flow_handler: FlowHandlerBase | None = None
 
-        self._schema_output_mapping: dict[str, type[BaseModel] | None] = {
+        self._schema_output_mapping: dict[str, type[BaseModel] | type[None]] = {
             WorkflowSchemaType.GENERATE: self._session_manager.workflow.single_output_schema,
             WorkflowSchemaType.CHAT: ChatResponse,
             WorkflowSchemaType.CHAT_STREAM: ChatResponseChunk,
@@ -105,45 +108,67 @@ class WebSocketMessageHandler:
                 if (isinstance(validated_message, WebSocketUserMessage)):
                     await self.process_workflow_request(validated_message)
 
-                elif isinstance(validated_message,
-                                (WebSocketSystemResponseTokenMessage,
-                                 WebSocketSystemIntermediateStepMessage,
-                                 WebSocketSystemInteractionMessage)):
+                elif isinstance(
+                        validated_message,
+                        WebSocketSystemResponseTokenMessage | WebSocketSystemIntermediateStepMessage
+                        | WebSocketSystemInteractionMessage):
                     # These messages are already handled by self.create_websocket_message(data_model=value, …)
                     # No further processing is needed here.
                     pass
 
                 elif (isinstance(validated_message, WebSocketUserInteractionResponseMessage)):
-                    user_content = await self.process_user_message_content(validated_message)
+                    user_content = await self._process_websocket_user_interaction_response_message(validated_message)
+                    assert self._user_interaction_response is not None
                     self._user_interaction_response.set_result(user_content)
             except (asyncio.CancelledError, WebSocketDisconnect):
                 # TODO: Handle the disconnect
                 break
 
-    async def process_user_message_content(
-            self, user_content: WebSocketUserMessage | WebSocketUserInteractionResponseMessage) -> BaseModel | None:
+    def _extract_last_user_message_content(self, messages: list[UserMessages]) -> TextContent:
         """
-        Processes the contents of a user message.
+        Extracts the last user's TextContent from a list of messages.
 
-        :param user_content: Incoming content data model.
-        :return: A validated Pydantic user content model or None if not found.
-        """
+        Args:
+            messages: List of UserMessages.
 
-        for user_message in user_content.content.messages[::-1]:
-            if (user_message.role == "user"):
+        Returns:
+            TextContent object from the last user message.
 
+        Raises:
+            ValueError: If no user text content is found.
+        """
+        for user_message in messages[::-1]:
+            if user_message.role == UserMessageContentRoleType.USER:
                 for attachment in user_message.content:
-
                     if isinstance(attachment, TextContent):
                         return attachment
+        raise ValueError("No user text content found in messages.")
+
+    async def _process_websocket_user_interaction_response_message(
+            self, user_content: WebSocketUserInteractionResponseMessage) -> TextContent:
+        """
+        Processes a WebSocketUserInteractionResponseMessage.
+        """
+        return self._extract_last_user_message_content(user_content.content.messages)
 
-        return None
+    async def _process_websocket_user_message(self, user_content: WebSocketUserMessage) -> ChatRequest | str:
+        """
+        Processes a WebSocketUserMessage based on schema type.
+        """
+        if self._workflow_schema_type in [WorkflowSchemaType.CHAT, WorkflowSchemaType.CHAT_STREAM]:
+            return ChatRequest(**user_content.content.model_dump(include={"messages"}))
+
+        elif self._workflow_schema_type in [WorkflowSchemaType.GENERATE, WorkflowSchemaType.GENERATE_STREAM]:
+            return self._extract_last_user_message_content(user_content.content.messages).text
+
+        raise ValueError("Unsupported workflow schema type for WebSocketUserMessage")
 
     async def process_workflow_request(self, user_message_as_validated_type: WebSocketUserMessage) -> None:
         """
         Process user messages and routes them appropriately.
 
-        :param user_message_as_validated_type: A WebSocketUserMessage Data Model instance.
+        Args:
+            user_message_as_validated_type (WebSocketUserMessage): The validated user message to process.
         """
 
         try:
@@ -151,18 +176,15 @@ class WebSocketMessageHandler:
             self._workflow_schema_type = user_message_as_validated_type.schema_type
             self._conversation_id = user_message_as_validated_type.conversation_id
 
-            content: BaseModel | None = await self.process_user_message_content(user_message_as_validated_type)
-
-            if content is None:
-                raise ValueError(f"User message content could not be found: {user_message_as_validated_type}")
+            message_content: typing.Any = await self._process_websocket_user_message(user_message_as_validated_type)
 
-            if isinstance(content, TextContent) and (self._running_workflow_task is None):
+            if (self._running_workflow_task is None):
 
-                def _done_callback(task: asyncio.Task):
+                def _done_callback(_task: asyncio.Task):
                     self._running_workflow_task = None
 
                 self._running_workflow_task = asyncio.create_task(
-                    self._run_workflow(payload=content.text,
+                    self._run_workflow(payload=message_content,
                                        user_message_id=self._message_parent_id,
                                        conversation_id=self._conversation_id,
                                        result_type=self._schema_output_mapping[self._workflow_schema_type],
@@ -180,13 +202,14 @@ class WebSocketMessageHandler:
     async def create_websocket_message(self,
                                        data_model: BaseModel,
                                        message_type: str | None = None,
-                                       status: str = WebSocketMessageStatus.IN_PROGRESS) -> None:
+                                       status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS) -> None:
         """
         Creates a websocket message that will be ready for routing based on message type or data model.
 
-        :param data_model: Message content model.
-        :param message_type: Message content model.
-        :param status: Message content model.
+        Args:
+            data_model (BaseModel): Message content model.
+            message_type (str | None): Message content model.
+            status (WebSocketMessageStatus): Message content model.
         """
         try:
             message: BaseModel | None = None
@@ -196,8 +219,8 @@ class WebSocketMessageHandler:
 
             message_schema: type[BaseModel] = await self._message_validator.get_message_schema_by_type(message_type)
 
-            if 'id' in data_model.model_fields:
-                message_id: str = data_model.id
+            if hasattr(data_model, 'id'):
+                message_id: str = str(getattr(data_model, 'id'))
             else:
                 message_id = str(uuid.uuid4())
 
@@ -253,12 +276,15 @@ class WebSocketMessageHandler:
         Registered human interaction callback that processes human interactions and returns
         responses from websocket connection.
 
-        :param prompt: Incoming interaction content data model.
-        :return: A Text Content Base Pydantic model.
+        Args:
+            prompt: Incoming interaction content data model.
+
+        Returns:
+            A Text Content Base Pydantic model.
         """
 
         # First create a future from the loop for the human response
-        human_response_future: asyncio.Future[HumanResponse] = asyncio.get_running_loop().create_future()
+        human_response_future: asyncio.Future[TextContent] = asyncio.get_running_loop().create_future()
 
         # Then add the future to the outstanding human prompts dictionary
         self._user_interaction_response = human_response_future
@@ -274,10 +300,10 @@ class WebSocketMessageHandler:
                 return HumanResponseNotification()
 
             # Wait for the human response future to complete
-            interaction_response: HumanResponse = await human_response_future
+            text_content: TextContent = await human_response_future
 
             interaction_response: HumanResponse = await self._message_validator.convert_text_content_to_human_response(
-                interaction_response, prompt.content)
+                text_content, prompt.content)
 
             return interaction_response
 
@@ -293,13 +319,12 @@ class WebSocketMessageHandler:
                             output_type: type | None = None) -> None:
 
         try:
-            async with self._session_manager.session(
-                    user_message_id=user_message_id,
-                    conversation_id=conversation_id,
-                    http_connection=self._socket,
-                    user_input_callback=self.human_interaction_callback,
-                    user_authentication_callback=(self._flow_handler.authenticate
-                                                  if self._flow_handler else None)) as session:
+            auth_callback = self._flow_handler.authenticate if self._flow_handler else None
+            async with self._session_manager.session(user_message_id=user_message_id,
+                                                     conversation_id=conversation_id,
+                                                     http_connection=self._socket,
+                                                     user_input_callback=self.human_interaction_callback,
+                                                     user_authentication_callback=auth_callback) as session:
 
                 async for value in generate_streaming_response(payload,
                                                                session_manager=session,

nat/front_ends/fastapi/message_validator.py

@@ -139,8 +139,10 @@ class MessageValidator:
                     text_content: str = str(data_model.payload)
                 validated_message_content = SystemResponseContent(text=text_content)
 
-            elif (isinstance(data_model, (ChatResponse, ChatResponseChunk))):
+            elif isinstance(data_model, ChatResponse):
                 validated_message_content = SystemResponseContent(text=data_model.choices[0].message.content)
+            elif isinstance(data_model, ChatResponseChunk):
+                validated_message_content = SystemResponseContent(text=data_model.choices[0].delta.content)
 
             elif (isinstance(data_model, ResponseIntermediateStep)):
                 validated_message_content = SystemIntermediateStepContent(name=data_model.name,
@@ -204,7 +206,7 @@ class MessageValidator:
 
         validated_message_type: str = ""
         try:
-            if (isinstance(data_model, (ResponsePayloadOutput, ChatResponse, ChatResponseChunk))):
+            if (isinstance(data_model, ResponsePayloadOutput | ChatResponse | ChatResponseChunk)):
                 validated_message_type = WebSocketMessageType.RESPONSE_MESSAGE
 
             elif (isinstance(data_model, ResponseIntermediateStep)):
@@ -238,10 +240,9 @@ class MessageValidator:
         thread_id: str = "default",
         parent_id: str = "default",
         conversation_id: str | None = None,
-        content: SystemResponseContent
-        | Error = SystemResponseContent(),
+        content: SystemResponseContent | Error = SystemResponseContent(),
         status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
-        timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+        timestamp: str = str(datetime.datetime.now(datetime.UTC))
     ) -> WebSocketSystemResponseTokenMessage | None:
         """
         Creates a system response token message with default values.
@@ -280,7 +281,7 @@ class MessageValidator:
         conversation_id: str | None = None,
         content: SystemIntermediateStepContent = SystemIntermediateStepContent(name="default", payload="default"),
         status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
-        timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+        timestamp: str = str(datetime.datetime.now(datetime.UTC))
     ) -> WebSocketSystemIntermediateStepMessage | None:
         """
         Creates a system intermediate step message with default values.
@@ -320,7 +321,7 @@ class MessageValidator:
         conversation_id: str | None = None,
         content: HumanPrompt,
         status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
-        timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+        timestamp: str = str(datetime.datetime.now(datetime.UTC))
     ) -> WebSocketSystemInteractionMessage | None:
         """
         Creates a system interaction message with default values.

nat/front_ends/fastapi/utils.py

@@ -0,0 +1,57 @@
+# 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 importlib
+import os
+
+
+def get_config_file_path() -> str:
+    """
+    Get the path to the NAT configuration file from the environment variable NAT_CONFIG_FILE.
+    Raises ValueError if the environment variable is not set.
+    """
+    config_file_path = os.getenv("NAT_CONFIG_FILE")
+    if (not config_file_path):
+        raise ValueError("Config file not found in environment variable NAT_CONFIG_FILE.")
+
+    return os.path.abspath(config_file_path)
+
+
+def import_class_from_string(class_full_name: str) -> type:
+    """
+    Import a class from a string in the format 'module.submodule.ClassName'.
+    Raises ImportError if the class cannot be imported.
+    """
+    try:
+        class_name_parts = class_full_name.split(".")
+
+        module_name = ".".join(class_name_parts[:-1])
+        class_name = class_name_parts[-1]
+
+        module = importlib.import_module(module_name)
+
+        if not hasattr(module, class_name):
+            raise ValueError(f"Class '{class_full_name}' not found.")
+
+        return getattr(module, class_name)
+    except (ImportError, AttributeError) as e:
+        raise ImportError(f"Could not import {class_full_name}.") from e
+
+
+def get_class_name(cls: type) -> str:
+    """
+    Get the full class name including the module.
+    """
+    return f"{cls.__module__}.{cls.__qualname__}"

nat/front_ends/mcp/introspection_token_verifier.py

@@ -0,0 +1,73 @@
+# 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.
+"""OAuth 2.0 Token Introspection verifier implementation for MCP servers."""
+
+import logging
+
+from mcp.server.auth.provider import AccessToken
+from mcp.server.auth.provider import TokenVerifier
+
+from nat.authentication.credential_validator.bearer_token_validator import BearerTokenValidator
+from nat.authentication.oauth2.oauth2_resource_server_config import OAuth2ResourceServerConfig
+
+logger = logging.getLogger(__name__)
+
+
+class IntrospectionTokenVerifier(TokenVerifier):
+    """Token verifier that delegates token verification to BearerTokenValidator."""
+
+    def __init__(self, config: OAuth2ResourceServerConfig):
+        """Create IntrospectionTokenVerifier from OAuth2ResourceServerConfig.
+
+        Args:
+            config: OAuth2ResourceServerConfig
+        """
+        issuer = config.issuer_url
+        scopes = config.scopes or []
+        audience = config.audience
+        jwks_uri = config.jwks_uri
+        introspection_endpoint = config.introspection_endpoint
+        discovery_url = config.discovery_url
+        client_id = config.client_id
+        client_secret = config.client_secret
+
+        self._bearer_token_validator = BearerTokenValidator(
+            issuer=issuer,
+            audience=audience,
+            scopes=scopes,
+            jwks_uri=jwks_uri,
+            introspection_endpoint=introspection_endpoint,
+            discovery_url=discovery_url,
+            client_id=client_id,
+            client_secret=client_secret,
+        )
+
+    async def verify_token(self, token: str) -> AccessToken | None:
+        """Verify token by delegating to BearerTokenValidator.
+
+        Args:
+            token: The Bearer token to verify
+
+        Returns:
+            AccessToken | None: AccessToken if valid, None if invalid
+        """
+        validation_result = await self._bearer_token_validator.verify(token)
+
+        if validation_result.active:
+            return AccessToken(token=token,
+                               expires_at=validation_result.expires_at,
+                               scopes=validation_result.scopes or [],
+                               client_id=validation_result.client_id or "")
+        return None

nat/front_ends/mcp/mcp_front_end_config.py

@@ -13,17 +13,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 from typing import Literal
 
 from pydantic import Field
+from pydantic import field_validator
+from pydantic import model_validator
 
+from nat.authentication.oauth2.oauth2_resource_server_config import OAuth2ResourceServerConfig
 from nat.data_models.front_end import FrontEndBaseConfig
 
+logger = logging.getLogger(__name__)
+
 
 class MCPFrontEndConfig(FrontEndBaseConfig, name="mcp"):
     """MCP front end configuration.
 
-    A simple MCP (Modular Communication Protocol) front end for NeMo Agent toolkit.
+    A simple MCP (Model Context Protocol) front end for NeMo Agent toolkit.
     """
 
     name: str = Field(default="NeMo Agent Toolkit MCP",
@@ -32,10 +38,72 @@ class MCPFrontEndConfig(FrontEndBaseConfig, name="mcp"):
     port: int = Field(default=9901, description="Port to bind the server to (default: 9901)", ge=0, le=65535)
     debug: bool = Field(default=False, description="Enable debug mode (default: False)")
     log_level: str = Field(default="INFO", description="Log level for the MCP server (default: INFO)")
-    tool_names: list[str] = Field(default_factory=list,
-                                  description="The list of tools MCP server will expose (default: all tools)")
+    tool_names: list[str] = Field(
+        default_factory=list,
+        description="The list of tools MCP server will expose (default: all tools)."
+        "Tool names can be functions or function groups",
+    )
     transport: Literal["sse", "streamable-http"] = Field(
         default="streamable-http",
         description="Transport type for the MCP server (default: streamable-http, backwards compatible with sse)")
     runner_class: str | None = Field(
         default=None, description="Custom worker class for handling MCP routes (default: built-in worker)")
+    base_path: str | None = Field(default=None,
+                                  description="Base path to mount the MCP server at (e.g., '/api/v1'). "
+                                  "If specified, the server will be accessible at http://host:port{base_path}/mcp. "
+                                  "If None, server runs at root path /mcp.")
+
+    server_auth: OAuth2ResourceServerConfig | None = Field(
+        default=None, description=("OAuth 2.0 Resource Server configuration for token verification."))
+
+    @field_validator('base_path')
+    @classmethod
+    def validate_base_path(cls, v: str | None) -> str | None:
+        """Validate that base_path starts with '/' and doesn't end with '/'."""
+        if v is not None:
+            if not v.startswith('/'):
+                raise ValueError("base_path must start with '/'")
+            if v.endswith('/'):
+                raise ValueError("base_path must not end with '/'")
+        return v
+
+    # Memory profiling configuration
+    enable_memory_profiling: bool = Field(default=False,
+                                          description="Enable memory profiling and diagnostics (default: False)")
+    memory_profile_interval: int = Field(default=50,
+                                         description="Log memory stats every N requests (default: 50)",
+                                         ge=1)
+    memory_profile_top_n: int = Field(default=10,
+                                      description="Number of top memory allocations to log (default: 10)",
+                                      ge=1,
+                                      le=50)
+    memory_profile_log_level: str = Field(default="DEBUG",
+                                          description="Log level for memory profiling output (default: DEBUG)")
+
+    @model_validator(mode="after")
+    def validate_security_configuration(self):
+        """Validate security configuration to prevent accidental misconfigurations."""
+        # Check if server is bound to a non-localhost interface without authentication
+        localhost_hosts = {"localhost", "127.0.0.1", "::1"}
+        if self.host not in localhost_hosts and self.server_auth is None:
+            logger.warning(
+                "MCP server is configured to bind to '%s' without authentication. "
+                "This may expose your server to unauthorized access. "
+                "Consider either: (1) binding to localhost for local-only access, "
+                "or (2) configuring server_auth for production deployments on public interfaces.",
+                self.host)
+
+        # Check if SSE transport is used (which doesn't support authentication)
+        if self.transport == "sse":
+            if self.server_auth is not None:
+                logger.warning("SSE transport does not support authentication. "
+                               "The configured server_auth will be ignored. "
+                               "For production use with authentication, use 'streamable-http' transport instead.")
+            elif self.host not in localhost_hosts:
+                logger.warning(
+                    "SSE transport does not support authentication and is bound to '%s'. "
+                    "This configuration is not recommended for production use. "
+                    "For production deployments, use 'streamable-http' transport with server_auth configured.",
+                    self.host)
+
+        return self