nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc1__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

@@ -86,7 +86,7 @@ class WebSocketMessageHandler:
 
     async def __aexit__(self, exc_type, exc_value, traceback) -> None:
 
-        # TODO: Handle the exit  # pylint: disable=fixme
+        # TODO: Handle the exit
         pass
 
     async def run(self) -> None:
@@ -107,10 +107,8 @@ class WebSocketMessageHandler:
 
                 elif isinstance(
                         validated_message,
-                    (  # noqa: E131
-                        WebSocketSystemResponseTokenMessage,
-                        WebSocketSystemIntermediateStepMessage,
-                        WebSocketSystemInteractionMessage)):
+                        WebSocketSystemResponseTokenMessage | WebSocketSystemIntermediateStepMessage
+                        | WebSocketSystemInteractionMessage):
                     # These messages are already handled by self.create_websocket_message(data_model=value, …)
                     # No further processing is needed here.
                     pass
@@ -119,11 +117,9 @@ class WebSocketMessageHandler:
                     user_content = await self.process_user_message_content(validated_message)
                     self._user_interaction_response.set_result(user_content)
             except (asyncio.CancelledError, WebSocketDisconnect):
-                # TODO: Handle the disconnect  # pylint: disable=fixme
+                # TODO: Handle the disconnect
                 break
 
-        return None
-
     async def process_user_message_content(
             self, user_content: WebSocketUserMessage | WebSocketUserInteractionResponseMessage) -> BaseModel | None:
         """
@@ -162,18 +158,19 @@ class WebSocketMessageHandler:
 
             if isinstance(content, TextContent) and (self._running_workflow_task is None):
 
-                def _done_callback(task: asyncio.Task):  # pylint: disable=unused-argument
+                def _done_callback(task: asyncio.Task):
                     self._running_workflow_task = None
 
                 self._running_workflow_task = asyncio.create_task(
-                    self._run_workflow(content.text,
-                                       self._conversation_id,
+                    self._run_workflow(payload=content.text,
+                                       user_message_id=self._message_parent_id,
+                                       conversation_id=self._conversation_id,
                                        result_type=self._schema_output_mapping[self._workflow_schema_type],
                                        output_type=self._schema_output_mapping[
                                            self._workflow_schema_type])).add_done_callback(_done_callback)
 
         except ValueError as e:
-            logger.error("User message content not found: %s", str(e), exc_info=True)
+            logger.exception("User message content not found: %s", str(e))
             await self.create_websocket_message(data_model=Error(code=ErrorTypes.INVALID_USER_MESSAGE_CONTENT,
                                                                  message="User message content could not be found",
                                                                  details=str(e)),
@@ -241,7 +238,7 @@ class WebSocketMessageHandler:
                     f"Message type could not be resolved by input data model: {data_model.model_dump_json()}")
 
         except (ValidationError, TypeError, ValueError) as e:
-            logger.error("A data vaidation error ocurred creating websocket message: %s", str(e), exc_info=True)
+            logger.exception("A data vaidation error ocurred creating websocket message: %s", str(e))
             message = await self._message_validator.create_system_response_token_message(
                 message_type=WebSocketMessageType.ERROR_MESSAGE,
                 conversation_id=self._conversation_id,
@@ -290,14 +287,16 @@ class WebSocketMessageHandler:
 
     async def _run_workflow(self,
                             payload: typing.Any,
+                            user_message_id: str | None = None,
                             conversation_id: str | None = None,
                             result_type: type | None = None,
                             output_type: type | None = None) -> None:
 
         try:
             async with self._session_manager.session(
+                    user_message_id=user_message_id,
                     conversation_id=conversation_id,
-                    request=self._socket,
+                    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:

nat/front_ends/fastapi/message_validator.py

@@ -97,7 +97,7 @@ class MessageValidator:
             return validated_message
 
         except (ValidationError, TypeError, ValueError) as e:
-            logger.error("A data validation error %s occurred for message: %s", str(e), str(message), exc_info=True)
+            logger.exception("A data validation error %s occurred for message: %s", str(e), str(message))
             return await self.create_system_response_token_message(message_type=WebSocketMessageType.ERROR_MESSAGE,
                                                                    content=Error(code=ErrorTypes.INVALID_MESSAGE,
                                                                                  message="Error validating message.",
@@ -119,7 +119,7 @@ class MessageValidator:
             return schema
 
         except (TypeError, ValueError) as e:
-            logger.error("Error retrieving schema for message type '%s': %s", message_type, str(e), exc_info=True)
+            logger.exception("Error retrieving schema for message type '%s': %s", message_type, str(e))
             return Error
 
     async def convert_data_to_message_content(self, data_model: BaseModel) -> BaseModel:
@@ -139,7 +139,7 @@ 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 | ChatResponseChunk)):
                 validated_message_content = SystemResponseContent(text=data_model.choices[0].message.content)
 
             elif (isinstance(data_model, ResponseIntermediateStep)):
@@ -156,7 +156,7 @@ class MessageValidator:
             return validated_message_content
 
         except ValueError as e:
-            logger.error("Input data could not be converted to validated message content: %s", str(e), exc_info=True)
+            logger.exception("Input data could not be converted to validated message content: %s", str(e))
             return Error(code=ErrorTypes.INVALID_DATA_CONTENT, message="Input data not supported.", details=str(e))
 
     async def convert_text_content_to_human_response(self, text_content: TextContent,
@@ -191,7 +191,7 @@ class MessageValidator:
             return human_response
 
         except ValueError as e:
-            logger.error("Error human response content not found: %s", str(e), exc_info=True)
+            logger.exception("Error human response content not found: %s", str(e))
             return HumanResponseText(text=str(e))
 
     async def resolve_message_type_by_data(self, data_model: BaseModel) -> str:
@@ -204,7 +204,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)):
@@ -218,9 +218,7 @@ class MessageValidator:
             return validated_message_type
 
         except ValueError as e:
-            logger.error("Error type not found converting data to validated websocket message content: %s",
-                         str(e),
-                         exc_info=True)
+            logger.exception("Error type not found converting data to validated websocket message content: %s", str(e))
             return WebSocketMessageType.ERROR_MESSAGE
 
     async def get_intermediate_step_parent_id(self, data_model: ResponseIntermediateStep) -> str:
@@ -232,7 +230,7 @@ class MessageValidator:
         """
         return data_model.parent_id or "root"
 
-    async def create_system_response_token_message(  # pylint: disable=R0917:too-many-positional-arguments
+    async def create_system_response_token_message(
         self,
         message_type: Literal[WebSocketMessageType.RESPONSE_MESSAGE,
                               WebSocketMessageType.ERROR_MESSAGE] = WebSocketMessageType.RESPONSE_MESSAGE,
@@ -243,7 +241,7 @@ class MessageValidator:
         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.
@@ -269,10 +267,10 @@ class MessageValidator:
                                                        timestamp=timestamp)
 
         except Exception as e:
-            logger.error("Error creating system response token message: %s", str(e), exc_info=True)
+            logger.exception("Error creating system response token message: %s", str(e))
             return None
 
-    async def create_system_intermediate_step_message(  # pylint: disable=R0917:too-many-positional-arguments
+    async def create_system_intermediate_step_message(
         self,
         message_type: Literal[WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE] = (
             WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE),
@@ -282,7 +280,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.
@@ -308,10 +306,10 @@ class MessageValidator:
                                                           timestamp=timestamp)
 
         except Exception as e:
-            logger.error("Error creating system intermediate step message: %s", str(e), exc_info=True)
+            logger.exception("Error creating system intermediate step message: %s", str(e))
             return None
 
-    async def create_system_interaction_message(  # pylint: disable=R0917:too-many-positional-arguments
+    async def create_system_interaction_message(
         self,
         *,
         message_type: Literal[WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE] = (
@@ -322,8 +320,8 @@ class MessageValidator:
         conversation_id: str | None = None,
         content: HumanPrompt,
         status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
-        timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
-    ) -> WebSocketSystemInteractionMessage | None:  # noqa: E125 continuation line with same indent as next logical line
+        timestamp: str = str(datetime.datetime.now(datetime.UTC))
+    ) -> WebSocketSystemInteractionMessage | None:
         """
         Creates a system interaction message with default values.
 
@@ -348,5 +346,5 @@ class MessageValidator:
                                                      timestamp=timestamp)
 
         except Exception as e:
-            logger.error("Error creating system interaction message: %s", str(e), exc_info=True)
+            logger.exception("Error creating system interaction message: %s", str(e))
             return None

nat/front_ends/fastapi/response_helpers.py

@@ -98,9 +98,9 @@ async def generate_streaming_response(payload: typing.Any,
                     yield item
                 else:
                     yield ResponsePayloadOutput(payload=item)
-        except Exception as e:
+        except Exception:
             # Handle exceptions here
-            raise e
+            raise
         finally:
             await q.close()
 
@@ -165,9 +165,9 @@ async def generate_streaming_response_full(payload: typing.Any,
                         yield item
                 else:
                     yield ResponsePayloadOutput(payload=item)
-        except Exception as e:
+        except Exception:
             # Handle exceptions here
-            raise e
+            raise
         finally:
             await q.close()
 

nat/front_ends/fastapi/step_adaptor.py

@@ -289,7 +289,7 @@ class StepAdaptor:
 
         return event
 
-    def process(self, step: IntermediateStep) -> ResponseSerializable | None:  # pylint: disable=R1710
+    def process(self, step: IntermediateStep) -> ResponseSerializable | None:
 
         # Track the chunk
         self._history.append(step)
@@ -314,6 +314,6 @@ class StepAdaptor:
                 return self._handle_custom(payload, ancestry)
 
         except Exception as e:
-            logger.error("Error processing intermediate step: %s", e, exc_info=True)
+            logger.exception("Error processing intermediate step: %s", e)
 
         return None

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,15 +13,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Literal
+
 from pydantic import Field
 
+from nat.authentication.oauth2.oauth2_resource_server_config import OAuth2ResourceServerConfig
 from nat.data_models.front_end import FrontEndBaseConfig
 
 
 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,5 +35,11 @@ class MCPFrontEndConfig(FrontEndBaseConfig, name="mcp"):
     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)")
+    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)")
+
+    server_auth: OAuth2ResourceServerConfig | None = Field(
+        default=None, description=("OAuth 2.0 Resource Server configuration for token verification."))

nat/front_ends/mcp/mcp_front_end_plugin.py

@@ -16,6 +16,7 @@
 import logging
 import typing
 
+from nat.authentication.oauth2.oauth2_resource_server_config import OAuth2ResourceServerConfig
 from nat.builder.front_end import FrontEndBase
 from nat.builder.workflow_builder import WorkflowBuilder
 from nat.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
@@ -55,27 +56,58 @@ class MCPFrontEndPlugin(FrontEndBase[MCPFrontEndConfig]):
 
         return worker_class(self.full_config)
 
+    async def _create_token_verifier(self, token_verifier_config: OAuth2ResourceServerConfig):
+        """Create a token verifier based on configuration."""
+        from nat.front_ends.mcp.introspection_token_verifier import IntrospectionTokenVerifier
+
+        if not self.front_end_config.server_auth:
+            return None
+
+        return IntrospectionTokenVerifier(token_verifier_config)
+
     async def run(self) -> None:
         """Run the MCP server."""
         # Import FastMCP
         from mcp.server.fastmcp import FastMCP
 
-        # Create an MCP server with the configured parameters
-        mcp = FastMCP(
-            self.front_end_config.name,
-            host=self.front_end_config.host,
-            port=self.front_end_config.port,
-            debug=self.front_end_config.debug,
-            log_level=self.front_end_config.log_level,
-        )
-
-        # Get the worker instance and set up routes
-        worker = self._get_worker_instance()
+        # Create auth settings and token verifier if auth is required
+        auth_settings = None
+        token_verifier = None
 
         # Build the workflow and add routes using the worker
         async with WorkflowBuilder.from_config(config=self.full_config) as builder:
+
+            if self.front_end_config.server_auth:
+                from mcp.server.auth.settings import AuthSettings
+                from pydantic import AnyHttpUrl
+
+                server_url = f"http://{self.front_end_config.host}:{self.front_end_config.port}"
+
+                auth_settings = AuthSettings(issuer_url=AnyHttpUrl(self.front_end_config.server_auth.issuer_url),
+                                             required_scopes=self.front_end_config.server_auth.scopes,
+                                             resource_server_url=AnyHttpUrl(server_url))
+
+                token_verifier = await self._create_token_verifier(self.front_end_config.server_auth)
+
+            # Create an MCP server with the configured parameters
+            mcp = FastMCP(name=self.front_end_config.name,
+                          host=self.front_end_config.host,
+                          port=self.front_end_config.port,
+                          debug=self.front_end_config.debug,
+                          auth=auth_settings,
+                          token_verifier=token_verifier)
+
+            # Get the worker instance and set up routes
+            worker = self._get_worker_instance()
+
             # Add routes through the worker (includes health endpoint and function registration)
             await worker.add_routes(mcp, builder)
 
-            # Start the MCP server
-            await mcp.run_sse_async()
+            # Start the MCP server with configurable transport
+            # streamable-http is the default, but users can choose sse if preferred
+            if self.front_end_config.transport == "sse":
+                logger.info("Starting MCP server with SSE endpoint at /sse")
+                await mcp.run_sse_async()
+            else:  # streamable-http
+                logger.info("Starting MCP server with streamable-http endpoint at /mcp/")
+                await mcp.run_streamable_http_async()