nvidia-nat 1.2.1rc1__py3-none-any.whl → 1.3.0__py3-none-any.whl

nat/front_ends/mcp/mcp_front_end_config.py

@@ -13,15 +13,22 @@
 # 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 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",
@@ -30,7 +37,44 @@ 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)")
+
+    server_auth: OAuth2ResourceServerConfig | None = Field(
+        default=None, description=("OAuth 2.0 Resource Server configuration for token verification."))
+
+    @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

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,61 @@ 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
+            try:
+                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()
+            except KeyboardInterrupt:
+                logger.info("MCP server shutdown requested (Ctrl+C). Shutting down gracefully.")

nat/front_ends/mcp/mcp_front_end_plugin_worker.py

@@ -16,11 +16,15 @@
 import logging
 from abc import ABC
 from abc import abstractmethod
+from collections.abc import Mapping
+from typing import Any
 
 from mcp.server.fastmcp import FastMCP
+from starlette.exceptions import HTTPException
 from starlette.requests import Request
 
 from nat.builder.function import Function
+from nat.builder.function_base import FunctionBase
 from nat.builder.workflow import Workflow
 from nat.builder.workflow_builder import WorkflowBuilder
 from nat.data_models.config import Config
@@ -82,7 +86,7 @@ class MCPFrontEndPluginWorkerBase(ABC):
         """
         pass
 
-    def _get_all_functions(self, workflow: Workflow) -> dict[str, Function]:
+    async def _get_all_functions(self, workflow: Workflow) -> dict[str, Function]:
         """Get all functions from the workflow.
 
         Args:
@@ -94,13 +98,114 @@ class MCPFrontEndPluginWorkerBase(ABC):
         functions: dict[str, Function] = {}
 
         # Extract all functions from the workflow
-        for function_name, function in workflow.functions.items():
-            functions[function_name] = function
+        functions.update(workflow.functions)
+        for function_group in workflow.function_groups.values():
+            functions.update(await function_group.get_accessible_functions())
 
-        functions[workflow.config.workflow.type] = workflow
+        if workflow.config.workflow.workflow_alias:
+            functions[workflow.config.workflow.workflow_alias] = workflow
+        else:
+            functions[workflow.config.workflow.type] = workflow
 
         return functions
 
+    def _setup_debug_endpoints(self, mcp: FastMCP, functions: Mapping[str, FunctionBase]) -> None:
+        """Set up HTTP debug endpoints for introspecting tools and schemas.
+
+        Exposes:
+          - GET /debug/tools/list: List tools. Optional query param `name` (one or more, repeatable or comma separated)
+            selects a subset and returns details for those tools.
+        """
+
+        @mcp.custom_route("/debug/tools/list", methods=["GET"])
+        async def list_tools(request: Request):
+            """HTTP list tools endpoint."""
+
+            from starlette.responses import JSONResponse
+
+            from nat.front_ends.mcp.tool_converter import get_function_description
+
+            # Query params
+            # Support repeated names and comma-separated lists
+            names_param_list = set(request.query_params.getlist("name"))
+            names: list[str] = []
+            for raw in names_param_list:
+                # if p.strip() is empty, it won't be included in the list!
+                parts = [p.strip() for p in raw.split(",") if p.strip()]
+                names.extend(parts)
+            detail_raw = request.query_params.get("detail")
+
+            def _parse_detail_param(detail_param: str | None, has_names: bool) -> bool:
+                if detail_param is None:
+                    if has_names:
+                        return True
+                    return False
+                v = detail_param.strip().lower()
+                if v in ("0", "false", "no", "off"):
+                    return False
+                if v in ("1", "true", "yes", "on"):
+                    return True
+                # For invalid values, default based on whether names are present
+                return has_names
+
+            # Helper function to build the input schema info
+            def _build_schema_info(fn: FunctionBase) -> dict[str, Any] | None:
+                schema = getattr(fn, "input_schema", None)
+                if schema is None:
+                    return None
+
+                # check if schema is a ChatRequest
+                schema_name = getattr(schema, "__name__", "")
+                schema_qualname = getattr(schema, "__qualname__", "")
+                if "ChatRequest" in schema_name or "ChatRequest" in schema_qualname:
+                    # Simplified interface used by MCP wrapper for ChatRequest
+                    return {
+                        "type": "object",
+                        "properties": {
+                            "query": {
+                                "type": "string", "description": "User query string"
+                            }
+                        },
+                        "required": ["query"],
+                        "title": "ChatRequestQuery",
+                    }
+
+                # Pydantic models provide model_json_schema
+                if schema is not None and hasattr(schema, "model_json_schema"):
+                    return schema.model_json_schema()
+
+                return None
+
+            def _build_final_json(functions_to_include: Mapping[str, FunctionBase],
+                                  include_schemas: bool = False) -> dict[str, Any]:
+                tools = []
+                for name, fn in functions_to_include.items():
+                    list_entry: dict[str, Any] = {
+                        "name": name, "description": get_function_description(fn), "is_workflow": hasattr(fn, "run")
+                    }
+                    if include_schemas:
+                        list_entry["schema"] = _build_schema_info(fn)
+                    tools.append(list_entry)
+
+                return {
+                    "count": len(tools),
+                    "tools": tools,
+                    "server_name": mcp.name,
+                }
+
+            if names:
+                # Return selected tools
+                try:
+                    functions_to_include = {n: functions[n] for n in names}
+                except KeyError as e:
+                    raise HTTPException(status_code=404, detail=f"Tool \"{e.args[0]}\" not found.") from e
+            else:
+                functions_to_include = functions
+
+            # Default for listing all: detail defaults to False unless explicitly set true
+            return JSONResponse(
+                _build_final_json(functions_to_include, _parse_detail_param(detail_raw, has_names=bool(names))))
+
 
 class MCPFrontEndPluginWorker(MCPFrontEndPluginWorkerBase):
     """Default MCP front end plugin worker implementation."""
@@ -118,10 +223,10 @@ class MCPFrontEndPluginWorker(MCPFrontEndPluginWorkerBase):
         self._setup_health_endpoint(mcp)
 
         # Build the workflow and register all functions with MCP
-        workflow = builder.build()
+        workflow = await builder.build()
 
         # Get all functions from the workflow
-        functions = self._get_all_functions(workflow)
+        functions = await self._get_all_functions(workflow)
 
         # Filter functions based on tool_names if provided
         if self.front_end_config.tool_names:
@@ -129,15 +234,22 @@ class MCPFrontEndPluginWorker(MCPFrontEndPluginWorkerBase):
             filtered_functions: dict[str, Function] = {}
             for function_name, function in functions.items():
                 if function_name in self.front_end_config.tool_names:
+                    # Treat current tool_names as function names, so check if the function name is in the list
+                    filtered_functions[function_name] = function
+                elif any(function_name.startswith(f"{group_name}.") for group_name in self.front_end_config.tool_names):
+                    # Treat tool_names as function group names, so check if the function name starts with the group name
                     filtered_functions[function_name] = function
                 else:
                     logger.debug("Skipping function %s as it's not in tool_names", function_name)
             functions = filtered_functions
 
-        # Register each function with MCP
+        # Register each function with MCP, passing workflow context for observability
         for function_name, function in functions.items():
-            register_function_with_mcp(mcp, function_name, function)
+            register_function_with_mcp(mcp, function_name, function, workflow)
 
         # Add a simple fallback function if no functions were found
         if not functions:
             raise RuntimeError("No functions found in workflow. Please check your configuration.")
+
+        # After registration, expose debug endpoints for tool/schema inspection
+        self._setup_debug_endpoints(mcp, functions)

nat/front_ends/mcp/tool_converter.py

@@ -17,13 +17,17 @@ import json
 import logging
 from inspect import Parameter
 from inspect import Signature
+from typing import TYPE_CHECKING
 
 from mcp.server.fastmcp import FastMCP
 from pydantic import BaseModel
 
+from nat.builder.context import ContextState
 from nat.builder.function import Function
 from nat.builder.function_base import FunctionBase
-from nat.builder.workflow import Workflow
+
+if TYPE_CHECKING:
+    from nat.builder.workflow import Workflow
 
 logger = logging.getLogger(__name__)
 
@@ -33,14 +37,16 @@ def create_function_wrapper(
     function: FunctionBase,
     schema: type[BaseModel],
     is_workflow: bool = False,
+    workflow: 'Workflow | None' = None,
 ):
     """Create a wrapper function that exposes the actual parameters of a NAT Function as an MCP tool.
 
     Args:
-        function_name: The name of the function/tool
-        function: The NAT Function object
-        schema: The input schema of the function
-        is_workflow: Whether the function is a Workflow
+        function_name (str): The name of the function/tool
+        function (FunctionBase): The NAT Function object
+        schema (type[BaseModel]): The input schema of the function
+        is_workflow (bool): Whether the function is a Workflow
+        workflow (Workflow | None): The parent workflow for observability context
 
     Returns:
         A wrapper function suitable for registration with MCP
@@ -101,6 +107,19 @@ def create_function_wrapper(
                 await ctx.report_progress(0, 100)
 
             try:
+                # Helper function to wrap function calls with observability
+                async def call_with_observability(func_call):
+                    # Use workflow's observability context (workflow should always be available)
+                    if not workflow:
+                        logger.error("Missing workflow context for function %s - observability will not be available",
+                                     function_name)
+                        raise RuntimeError("Workflow context is required for observability")
+
+                    logger.debug("Starting observability context for function %s", function_name)
+                    context_state = ContextState.get()
+                    async with workflow.exporter_manager.start(context_state=context_state):
+                        return await func_call()
+
                 # Special handling for ChatRequest
                 if is_chat_request:
                     from nat.data_models.api_server import ChatRequest
@@ -118,7 +137,7 @@ def create_function_wrapper(
                             result = await runner.result(to_type=str)
                     else:
                         # Regular functions use ainvoke
-                        result = await function.ainvoke(chat_request, to_type=str)
+                        result = await call_with_observability(lambda: function.ainvoke(chat_request, to_type=str))
                 else:
                     # Regular handling
                     # Handle complex input schema - if we extracted fields from a nested schema,
@@ -129,7 +148,7 @@ def create_function_wrapper(
                         field_type = schema.model_fields[field_name].annotation
 
                         # If it's a pydantic model, we need to create an instance
-                        if hasattr(field_type, "model_validate"):
+                        if field_type and hasattr(field_type, "model_validate"):
                             # Create the nested object
                             nested_obj = field_type.model_validate(kwargs)
                             # Call with the nested object
@@ -147,7 +166,7 @@ def create_function_wrapper(
                             result = await runner.result(to_type=str)
                     else:
                         # Regular function call
-                        result = await function.acall_invoke(**kwargs)
+                        result = await call_with_observability(lambda: function.acall_invoke(**kwargs))
 
                 # Report completion
                 if ctx:
@@ -156,7 +175,7 @@ def create_function_wrapper(
                 # Handle different result types for proper formatting
                 if isinstance(result, str):
                     return result
-                if isinstance(result, (dict, list)):
+                if isinstance(result, dict | list):
                     return json.dumps(result, default=str)
                 return str(result)
             except Exception as e:
@@ -170,7 +189,7 @@ def create_function_wrapper(
     wrapper = create_wrapper()
 
     # Set the signature on the wrapper function (WITHOUT ctx)
-    wrapper.__signature__ = sig
+    wrapper.__signature__ = sig  # type: ignore
     wrapper.__name__ = function_name
 
     # Return the wrapper with proper signature
@@ -183,8 +202,8 @@ def get_function_description(function: FunctionBase) -> str:
 
     The description is determined using the following precedence:
        1. If the function is a Workflow and has a 'description' attribute, use it.
-       2. If the Workflow's config has a 'topic', use it.
-       3. If the Workflow's config has a 'description', use it.
+       2. If the Workflow's config has a 'description', use it.
+       3. If the Workflow's config has a 'topic', use it.
        4. If the function is a regular Function, use its 'description' attribute.
 
     Args:
@@ -195,6 +214,9 @@ def get_function_description(function: FunctionBase) -> str:
     """
     function_description = ""
 
+    # Import here to avoid circular imports
+    from nat.builder.workflow import Workflow
+
     if isinstance(function, Workflow):
         config = function.config
 
@@ -207,6 +229,9 @@ def get_function_description(function: FunctionBase) -> str:
         # Try to get anything that might be a description
         elif hasattr(config, "topic") and config.topic:
             function_description = config.topic
+        # Try to get description from the workflow config
+        elif hasattr(config, "workflow") and hasattr(config.workflow, "description") and config.workflow.description:
+            function_description = config.workflow.description
 
     elif isinstance(function, Function):
         function_description = function.description
@@ -214,13 +239,17 @@ def get_function_description(function: FunctionBase) -> str:
     return function_description
 
 
-def register_function_with_mcp(mcp: FastMCP, function_name: str, function: FunctionBase) -> None:
+def register_function_with_mcp(mcp: FastMCP,
+                               function_name: str,
+                               function: FunctionBase,
+                               workflow: 'Workflow | None' = None) -> None:
     """Register a NAT Function as an MCP tool.
 
     Args:
         mcp: The FastMCP instance
         function_name: The name to register the function under
         function: The NAT Function to register
+        workflow: The parent workflow for observability context (if available)
     """
     logger.info("Registering function %s with MCP", function_name)
 
@@ -229,6 +258,7 @@ def register_function_with_mcp(mcp: FastMCP, function_name: str, function: Funct
     logger.info("Function %s has input schema: %s", function_name, input_schema)
 
     # Check if we're dealing with a Workflow
+    from nat.builder.workflow import Workflow
     is_workflow = isinstance(function, Workflow)
     if is_workflow:
         logger.info("Function %s is a Workflow", function_name)
@@ -237,5 +267,5 @@ def register_function_with_mcp(mcp: FastMCP, function_name: str, function: Funct
     function_description = get_function_description(function)
 
     # Create and register the wrapper function with MCP
-    wrapper_func = create_function_wrapper(function_name, function, input_schema, is_workflow)
+    wrapper_func = create_function_wrapper(function_name, function, input_schema, is_workflow, workflow)
     mcp.tool(name=function_name, description=function_description)(wrapper_func)

nat/front_ends/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
 

nat/front_ends/simple_base/simple_front_end_plugin_base.py

@@ -35,6 +35,8 @@ class SimpleFrontEndPluginBase(FrontEndBase[FrontEndConfigT], ABC):
 
     async def run(self):
 
+        await self.pre_run()
+
         # Must yield the workflow function otherwise it cleans up
         async with WorkflowBuilder.from_config(config=self.full_config) as builder:
 
@@ -45,7 +47,7 @@ class SimpleFrontEndPluginBase(FrontEndBase[FrontEndConfigT], ABC):
 
                 click.echo(stream.getvalue())
 
-            workflow = builder.build()
+            workflow = await builder.build()
             session_manager = SessionManager(workflow)
             await self.run_workflow(session_manager)
 

nat/llm/aws_bedrock_llm.py

@@ -21,27 +21,39 @@ from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
 from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import OptimizableMixin
+from nat.data_models.optimizable import SearchSpace
 from nat.data_models.retry_mixin import RetryMixin
+from nat.data_models.temperature_mixin import TemperatureMixin
+from nat.data_models.thinking_mixin import ThinkingMixin
+from nat.data_models.top_p_mixin import TopPMixin
 
 
-class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, name="aws_bedrock"):
+class AWSBedrockModelConfig(LLMBaseConfig,
+                            RetryMixin,
+                            OptimizableMixin,
+                            TemperatureMixin,
+                            TopPMixin,
+                            ThinkingMixin,
+                            name="aws_bedrock"):
     """An AWS Bedrock llm provider to be used with an LLM client."""
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
     # Completion parameters
     model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
                             serialization_alias="model",
                             description="The model name for the hosted AWS Bedrock.")
-    temperature: float = Field(default=0.0, ge=0.0, le=1.0, description="Sampling temperature in [0, 1].")
-    max_tokens: int | None = Field(default=1024,
-                                   gt=0,
-                                   description="Maximum number of tokens to generate."
-                                   "This field is ONLY required when using AWS Bedrock with Langchain.")
-    context_size: int | None = Field(default=1024,
-                                     gt=0,
-                                     description="Maximum number of tokens to generate."
-                                     "This field is ONLY required when using AWS Bedrock with LlamaIndex.")
+    max_tokens: int = OptimizableField(default=300,
+                                       description="Maximum number of tokens to generate.",
+                                       space=SearchSpace(high=2176, low=128, step=512))
+    context_size: int | None = Field(
+        default=1024,
+        gt=0,
+        description="The maximum number of tokens available for input. This is only required for LlamaIndex. "
+        "This field is ignored for LangChain/LangGraph.",
+    )
 
     # Client parameters
     region_name: str | None = Field(default="None", description="AWS region to use.")
@@ -52,6 +64,6 @@ class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, name="aws_bedrock"):
 
 
 @register_llm_provider(config_type=AWSBedrockModelConfig)
-async def aws_bedrock_model(llm_config: AWSBedrockModelConfig, builder: Builder):
+async def aws_bedrock_model(llm_config: AWSBedrockModelConfig, _builder: Builder):
 
     yield LLMProviderInfo(config=llm_config, description="A AWS Bedrock model for use with an LLM client.")

nat/llm/azure_openai_llm.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.
+
+from pydantic import AliasChoices
+from pydantic import ConfigDict
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.llm import LLMProviderInfo
+from nat.cli.register_workflow import register_llm_provider
+from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.retry_mixin import RetryMixin
+from nat.data_models.temperature_mixin import TemperatureMixin
+from nat.data_models.thinking_mixin import ThinkingMixin
+from nat.data_models.top_p_mixin import TopPMixin
+
+
+class AzureOpenAIModelConfig(
+        LLMBaseConfig,
+        RetryMixin,
+        TemperatureMixin,
+        TopPMixin,
+        ThinkingMixin,
+        name="azure_openai",
+):
+    """An Azure OpenAI LLM provider to be used with an LLM client."""
+
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
+
+    api_key: str | None = Field(default=None, description="Azure OpenAI API key to interact with hosted model.")
+    api_version: str = Field(default="2025-04-01-preview", description="Azure OpenAI API version.")
+    azure_endpoint: str | None = Field(validation_alias=AliasChoices("azure_endpoint", "base_url"),
+                                       serialization_alias="azure_endpoint",
+                                       default=None,
+                                       description="Base URL for the hosted model.")
+    azure_deployment: str = Field(validation_alias=AliasChoices("azure_deployment", "model_name", "model"),
+                                  serialization_alias="azure_deployment",
+                                  description="The Azure OpenAI hosted model/deployment name.")
+    seed: int | None = Field(default=None, description="Random seed to set for generation.")
+
+
+@register_llm_provider(config_type=AzureOpenAIModelConfig)
+async def azure_openai_llm(config: AzureOpenAIModelConfig, _builder: Builder):
+
+    yield LLMProviderInfo(config=config, description="An Azure OpenAI model for use with an LLM client.")