nvidia-nat 1.4.0a20251112__py3-none-any.whl → 1.4.0a20251120__py3-none-any.whl

nat/builder/builder.py

@@ -32,6 +32,7 @@ from nat.data_models.component_ref import FunctionGroupRef
 from nat.data_models.component_ref import FunctionRef
 from nat.data_models.component_ref import LLMRef
 from nat.data_models.component_ref import MemoryRef
+from nat.data_models.component_ref import MiddlewareRef
 from nat.data_models.component_ref import ObjectStoreRef
 from nat.data_models.component_ref import RetrieverRef
 from nat.data_models.component_ref import TTCStrategyRef
@@ -42,6 +43,7 @@ from nat.data_models.function import FunctionGroupBaseConfig
 from nat.data_models.function_dependencies import FunctionDependencies
 from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.memory import MemoryBaseConfig
+from nat.data_models.middleware import MiddlewareBaseConfig
 from nat.data_models.object_store import ObjectStoreBaseConfig
 from nat.data_models.retriever import RetrieverBaseConfig
 from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
@@ -49,6 +51,7 @@ from nat.experimental.decorators.experimental_warning_decorator import experimen
 from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
 from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
 from nat.memory.interfaces import MemoryEditor
+from nat.middleware.middleware import Middleware
 from nat.object_store.interfaces import ObjectStore
 from nat.retriever.interface import Retriever
 
@@ -289,6 +292,55 @@ class Builder(ABC):
     def get_function_group_dependencies(self, fn_name: str) -> FunctionDependencies:
         pass
 
+    @abstractmethod
+    async def add_middleware(self, name: str | MiddlewareRef, config: MiddlewareBaseConfig) -> Middleware:
+        """Add middleware to the builder.
+
+        Args:
+            name: The name or reference for the middleware
+            config: The configuration for the middleware
+
+        Returns:
+            The built middleware instance
+        """
+        pass
+
+    @abstractmethod
+    async def get_middleware(self, middleware_name: str | MiddlewareRef) -> Middleware:
+        """Get built middleware by name.
+
+        Args:
+            middleware_name: The name or reference of the middleware
+
+        Returns:
+            The built middleware instance
+        """
+        pass
+
+    @abstractmethod
+    def get_middleware_config(self, middleware_name: str | MiddlewareRef) -> MiddlewareBaseConfig:
+        """Get the configuration for middleware.
+
+        Args:
+            middleware_name: The name or reference of the middleware
+
+        Returns:
+            The configuration for the middleware
+        """
+        pass
+
+    async def get_middleware_list(self, middleware_names: Sequence[str | MiddlewareRef]) -> list[Middleware]:
+        """Get multiple middleware by name.
+
+        Args:
+            middleware_names: The names or references of the middleware
+
+        Returns:
+            List of built middleware instances
+        """
+        tasks = [self.get_middleware(name) for name in middleware_names]
+        return list(await asyncio.gather(*tasks, return_exceptions=False))
+
 
 class EvalBuilder(ABC):
 

nat/builder/component_utils.py

@@ -33,6 +33,7 @@ from nat.data_models.function import FunctionBaseConfig
 from nat.data_models.function import FunctionGroupBaseConfig
 from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.memory import MemoryBaseConfig
+from nat.data_models.middleware import MiddlewareBaseConfig
 from nat.data_models.object_store import ObjectStoreBaseConfig
 from nat.data_models.retriever import RetrieverBaseConfig
 from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
@@ -41,6 +42,7 @@ from nat.utils.type_utils import DecomposedType
 logger = logging.getLogger(__name__)
 
 # Order in which we want to process the component groups
+# IMPORTANT: MIDDLEWARE must be built before FUNCTIONS
 _component_group_order = [
     ComponentGroup.AUTHENTICATION,
     ComponentGroup.EMBEDDERS,
@@ -49,6 +51,7 @@ _component_group_order = [
     ComponentGroup.OBJECT_STORES,
     ComponentGroup.RETRIEVERS,
     ComponentGroup.TTC_STRATEGIES,
+    ComponentGroup.MIDDLEWARE,
     ComponentGroup.FUNCTION_GROUPS,
     ComponentGroup.FUNCTIONS,
 ]
@@ -111,6 +114,8 @@ def group_from_component(component: TypedBaseModel) -> ComponentGroup | None:
         return ComponentGroup.FUNCTIONS
     if (isinstance(component, FunctionGroupBaseConfig)):
         return ComponentGroup.FUNCTION_GROUPS
+    if (isinstance(component, MiddlewareBaseConfig)):
+        return ComponentGroup.MIDDLEWARE
     if (isinstance(component, LLMBaseConfig)):
         return ComponentGroup.LLMS
     if (isinstance(component, MemoryBaseConfig)):
@@ -260,7 +265,8 @@ def build_dependency_sequence(config: "Config") -> list[ComponentInstanceData]:
 
     total_node_count = (len(config.embedders) + len(config.functions) + len(config.function_groups) + len(config.llms) +
                         len(config.memory) + len(config.object_stores) + len(config.retrievers) +
-                        len(config.ttc_strategies) + len(config.authentication) + 1)  # +1 for the workflow
+                        len(config.ttc_strategies) + len(config.authentication) + len(config.middleware) + 1
+                        )  # +1 for the workflow
 
     dependency_map: dict
     dependency_graph: nx.DiGraph

nat/builder/context.py

@@ -34,6 +34,7 @@ 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.invocation_node import InvocationNode
+from nat.data_models.runtime_enum import RuntimeTypeEnum
 from nat.runtime.user_metadata import RequestAttributes
 from nat.utils.reactive.subject import Subject
 
@@ -72,6 +73,8 @@ class ContextState(metaclass=Singleton):
         self.workflow_trace_id: ContextVar[int | None] = ContextVar("workflow_trace_id", default=None)
         self.input_message: ContextVar[typing.Any] = ContextVar("input_message", default=None)
         self.user_manager: ContextVar[typing.Any] = ContextVar("user_manager", default=None)
+        self.runtime_type: ContextVar[RuntimeTypeEnum] = ContextVar("runtime_type",
+                                                                    default=RuntimeTypeEnum.RUN_OR_SERVE)
         self._metadata: ContextVar[RequestAttributes | None] = ContextVar("request_attributes", default=None)
         self._event_stream: ContextVar[Subject[IntermediateStep] | None] = ContextVar("event_stream", default=None)
         self._active_function: ContextVar[InvocationNode | None] = ContextVar("active_function", default=None)
@@ -302,6 +305,20 @@ class Context:
             raise RuntimeError("User authentication callback is not set in the context.")
         return callback
 
+    @property
+    def is_evaluating(self) -> bool:
+        """
+        Indicates whether the current context is in evaluation mode.
+
+        This property checks the context state to determine if the current
+        operation is being performed in evaluation mode. It returns a boolean
+        value indicating the evaluation status.
+
+        Returns:
+            bool: True if in evaluation mode, False otherwise.
+        """
+        return self._context_state.runtime_type.get() == RuntimeTypeEnum.EVALUATE
+
     @staticmethod
     def get() -> "Context":
         """

nat/builder/framework_enum.py

@@ -23,3 +23,4 @@ class LLMFrameworkEnum(str, Enum):
     SEMANTIC_KERNEL = "semantic_kernel"
     AGNO = "agno"
     ADK = "adk"
+    STRANDS = "strands"

nat/builder/function.py

@@ -34,6 +34,9 @@ from nat.builder.function_info import FunctionInfo
 from nat.data_models.function import EmptyFunctionConfig
 from nat.data_models.function import FunctionBaseConfig
 from nat.data_models.function import FunctionGroupBaseConfig
+from nat.middleware.function_middleware import FunctionMiddlewareChain
+from nat.middleware.middleware import FunctionMiddlewareContext
+from nat.middleware.middleware import Middleware
 
 _InvokeFnT = Callable[[InputT], Awaitable[SingleOutputT]]
 _StreamFnT = Callable[[InputT], AsyncGenerator[StreamingOutputT]]
@@ -64,6 +67,9 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
         self.description = description
         self.instance_name = instance_name or config.type
         self._context = Context.get()
+        self._configured_middleware: tuple[Middleware, ...] = tuple()
+        self._middlewared_single: _InvokeFnT | None = None
+        self._middlewared_stream: _StreamFnT | None = None
 
     def convert(self, value: typing.Any, to_type: type[_T]) -> _T:
         """
@@ -108,6 +114,38 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
         """
         return self._converter.try_convert(value, to_type=to_type)
 
+    @property
+    def middleware(self) -> tuple[Middleware, ...]:
+        """Return the currently configured middleware chain."""
+
+        return self._configured_middleware
+
+    def configure_middleware(self, middleware: Sequence[Middleware] | None = None) -> None:
+        """Attach an ordered list of middleware to this function instance."""
+
+        middleware_tuple: tuple[Middleware, ...] = tuple(middleware or ())
+
+        self._configured_middleware = middleware_tuple
+
+        if not middleware_tuple:
+            self._middlewared_single = None
+            self._middlewared_stream = None
+            return
+
+        logger.info(f"Building middleware for function '{self.instance_name}' in order of: {middleware_tuple}")
+
+        context = FunctionMiddlewareContext(name=self.instance_name,
+                                            config=self.config,
+                                            description=self.description,
+                                            input_schema=self.input_schema,
+                                            single_output_schema=self.single_output_schema,
+                                            stream_output_schema=self.streaming_output_schema)
+
+        chain = FunctionMiddlewareChain(middleware=middleware_tuple, context=context)
+
+        self._middlewared_single = chain.build_single(self._ainvoke) if self.has_single_output else None
+        self._middlewared_stream = chain.build_stream(self._astream) if self.has_streaming_output else None
+
     @abstractmethod
     async def _ainvoke(self, value: InputT) -> SingleOutputT:
         pass
@@ -150,7 +188,9 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
             try:
                 converted_input: InputT = self._convert_input(value)
 
-                result = await self._ainvoke(converted_input)
+                invoke_callable = self._middlewared_single or self._ainvoke
+
+                result = await invoke_callable(converted_input)
 
                 if to_type is not None and not isinstance(result, to_type):
                     result = self.convert(result, to_type)
@@ -243,7 +283,9 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
                 # Collect streaming outputs to capture the final result
                 final_output: list[typing.Any] = []
 
-                async for data in self._astream(converted_input):
+                stream_callable = self._middlewared_stream or self._astream
+
+                async for data in stream_callable(converted_input):
                     if to_type is not None and not isinstance(data, to_type):
                         converted_data = self.convert(data, to_type=to_type)
                         final_output.append(converted_data)
@@ -357,7 +399,8 @@ class FunctionGroup:
                  *,
                  config: FunctionGroupBaseConfig,
                  instance_name: str | None = None,
-                 filter_fn: Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None = None):
+                 filter_fn: Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None = None,
+                 middleware: Sequence[Middleware] | None = None):
         """
         Creates a new function group.
 
@@ -370,12 +413,15 @@ class FunctionGroup:
         filter_fn : Callable[[Sequence[str]], Awaitable[Sequence[str]]] | None, optional
             A callback function to additionally filter the functions in the function group dynamically when
             the functions are accessed via any accessor method.
+        middleware : Sequence[Middleware] | None, optional
+            The middleware instances to apply to all functions in this group.
         """
         self._config = config
         self._instance_name = instance_name or config.type
         self._functions: dict[str, Function] = dict()
         self._filter_fn = filter_fn
         self._per_function_filter_fn: dict[str, Callable[[str], Awaitable[bool]]] = dict()
+        self._middleware: tuple[Middleware, ...] = tuple(middleware or ())
 
     def add_function(self,
                      name: str,
@@ -424,6 +470,9 @@ class FunctionGroup:
         info = FunctionInfo.from_fn(fn, input_schema=input_schema, description=description, converters=converters)
         full_name = self._get_fn_name(name)
         lambda_fn = LambdaFunction.from_info(config=EmptyFunctionConfig(), info=info, instance_name=full_name)
+        # Configure middleware from the function group if any
+        if self._middleware:
+            lambda_fn.configure_middleware(self._middleware)
         self._functions[name] = lambda_fn
         if filter_fn:
             self._per_function_filter_fn[name] = filter_fn
@@ -712,3 +761,25 @@ class FunctionGroup:
         Returns the instance name for the function group.
         """
         return self._instance_name
+
+    @property
+    def middleware(self) -> tuple[Middleware, ...]:
+        """
+        Returns the middleware configured for this function group.
+        """
+        return self._middleware
+
+    def configure_middleware(self, middleware: Sequence[Middleware] | None = None) -> None:
+        """
+        Configure the middleware for this function group.
+        These middleware will be applied to all functions added to the group.
+
+        Parameters
+        ----------
+        middleware : Sequence[Middleware] | None
+            The middleware to configure for the function group.
+        """
+        self._middleware = tuple(middleware or ())
+        # Update existing functions with the new middleware
+        for func in self._functions.values():
+            func.configure_middleware(self._middleware)

nat/builder/workflow.py

@@ -28,6 +28,7 @@ from nat.builder.function_base import StreamingOutputT
 from nat.builder.llm import LLMProviderInfo
 from nat.builder.retriever import RetrieverProviderInfo
 from nat.data_models.config import Config
+from nat.data_models.runtime_enum import RuntimeTypeEnum
 from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
 from nat.memory.interfaces import MemoryEditor
 from nat.object_store.interfaces import ObjectStore
@@ -94,7 +95,7 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
         return self._exporter_manager.get()
 
     @asynccontextmanager
-    async def run(self, message: InputT):
+    async def run(self, message: InputT, runtime_type: RuntimeTypeEnum = RuntimeTypeEnum.RUN_OR_SERVE):
         """
         Called each time we start a new workflow run. We'll create
         a new top-level workflow span here.
@@ -103,7 +104,8 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
         async with Runner(input_message=message,
                           entry_fn=self._entry_fn,
                           context_state=self._context_state,
-                          exporter_manager=self.exporter_manager) as runner:
+                          exporter_manager=self.exporter_manager,
+                          runtime_type=runtime_type) as runner:
 
             # The caller can `yield runner` so they can do `runner.result()` or `runner.result_stream()`
             yield runner

nat/builder/workflow_builder.py

@@ -51,6 +51,7 @@ from nat.data_models.component_ref import FunctionGroupRef
 from nat.data_models.component_ref import FunctionRef
 from nat.data_models.component_ref import LLMRef
 from nat.data_models.component_ref import MemoryRef
+from nat.data_models.component_ref import MiddlewareRef
 from nat.data_models.component_ref import ObjectStoreRef
 from nat.data_models.component_ref import RetrieverRef
 from nat.data_models.component_ref import TTCStrategyRef
@@ -62,6 +63,7 @@ from nat.data_models.function import FunctionGroupBaseConfig
 from nat.data_models.function_dependencies import FunctionDependencies
 from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.memory import MemoryBaseConfig
+from nat.data_models.middleware import MiddlewareBaseConfig
 from nat.data_models.object_store import ObjectStoreBaseConfig
 from nat.data_models.retriever import RetrieverBaseConfig
 from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig
@@ -71,6 +73,8 @@ from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEn
 from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
 from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
 from nat.memory.interfaces import MemoryEditor
+from nat.middleware.function_middleware import FunctionMiddleware
+from nat.middleware.middleware import Middleware
 from nat.object_store.interfaces import ObjectStore
 from nat.observability.exporter.base_exporter import BaseExporter
 from nat.profiler.decorators.framework_wrapper import chain_wrapped_build_fn
@@ -141,6 +145,12 @@ class ConfiguredTTCStrategy:
     instance: StrategyBase
 
 
+@dataclasses.dataclass
+class ConfiguredMiddleware:
+    config: MiddlewareBaseConfig
+    instance: Middleware
+
+
 class WorkflowBuilder(Builder, AbstractAsyncContextManager):
 
     def __init__(self, *, general_config: GeneralConfig | None = None, registry: TypeRegistry | None = None):
@@ -170,6 +180,7 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
         self._object_stores: dict[str, ConfiguredObjectStore] = {}
         self._retrievers: dict[str, ConfiguredRetriever] = {}
         self._ttc_strategies: dict[str, ConfiguredTTCStrategy] = {}
+        self._middleware: dict[str, ConfiguredMiddleware] = {}
 
         self._context_state = ContextState.get()
 
@@ -423,6 +434,22 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
             raise ValueError("Expected a function, FunctionInfo object, or FunctionBase object to be "
                              f"returned from the function builder. Got {type(build_result)}")
 
+        # Resolve middleware names from config to middleware instances
+        # Only FunctionMiddleware types can be used with functions
+        middleware_instances = []
+        for middleware_name in config.middleware:
+            if middleware_name not in self._middleware:
+                raise ValueError(f"Middleware `{middleware_name}` not found for function `{name}`. "
+                                 f"It must be configured in the `middleware` section of the YAML configuration.")
+            middleware_obj = self._middleware[middleware_name].instance
+            if not isinstance(middleware_obj, FunctionMiddleware):
+                raise TypeError(
+                    f"Middleware `{middleware_name}` is not a FunctionMiddleware and cannot be used with functions. "
+                    f"Only FunctionMiddleware types support function-specific wrapping.")
+            middleware_instances.append(middleware_obj)
+
+        build_result.configure_middleware(middleware_instances)
+
         return ConfiguredFunction(config=config, instance=build_result)
 
     async def _build_function_group(self, name: str, config: FunctionGroupBaseConfig) -> ConfiguredFunctionGroup:
@@ -461,6 +488,23 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
             raise ValueError("Expected a FunctionGroup object to be returned from the function group builder. "
                              f"Got {type(build_result)}")
 
+        # Resolve middleware names from config to middleware instances
+        # Only FunctionMiddleware types can be used with function groups
+        middleware_instances = []
+        for middleware_name in config.middleware:
+            if middleware_name not in self._middleware:
+                raise ValueError(f"Middleware `{middleware_name}` not found for function group `{name}`. "
+                                 f"It must be configured in the `middleware` section of the YAML configuration.")
+            middleware_obj = self._middleware[middleware_name].instance
+            if not isinstance(middleware_obj, FunctionMiddleware):
+                raise TypeError(f"Middleware `{middleware_name}` is not a FunctionMiddleware and "
+                                f"cannot be used with function groups. "
+                                f"Only FunctionMiddleware types support function-specific wrapping.")
+            middleware_instances.append(middleware_obj)
+
+        # Configure middleware for the function group
+        build_result.configure_middleware(middleware_instances)
+
         # set the instance name for the function group based on the workflow-provided name
         build_result.set_instance_name(name)
         return ConfiguredFunctionGroup(config=config, instance=build_result)
@@ -968,6 +1012,72 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
 
         return config
 
+    @override
+    async def add_middleware(self, name: str | MiddlewareRef, config: MiddlewareBaseConfig) -> Middleware:
+        """Add middleware to the builder.
+
+        Args:
+            name: The name or reference for the middleware
+            config: The configuration for the middleware
+
+        Returns:
+            The built middleware instance
+
+        Raises:
+            ValueError: If the middleware already exists
+        """
+        if name in self._middleware:
+            raise ValueError(f"Middleware `{name}` already exists in the list of middleware")
+
+        try:
+            middleware_info = self._registry.get_middleware(type(config))
+
+            middleware_instance = await self._get_exit_stack().enter_async_context(
+                middleware_info.build_fn(config, self))
+
+            self._middleware[name] = ConfiguredMiddleware(config=config, instance=middleware_instance)
+
+            return middleware_instance
+        except Exception as e:
+            logger.error("Error adding function middleware `%s` with config `%s`: %s", name, config, e)
+            raise
+
+    @override
+    async def get_middleware(self, middleware_name: str | MiddlewareRef) -> Middleware:
+        """Get built middleware by name.
+
+        Args:
+            middleware_name: The name or reference of the middleware
+
+        Returns:
+            The built middleware instance
+
+        Raises:
+            ValueError: If the middleware is not found
+        """
+        if middleware_name not in self._middleware:
+            raise ValueError(f"Middleware `{middleware_name}` not found")
+
+        return self._middleware[middleware_name].instance
+
+    @override
+    def get_middleware_config(self, middleware_name: str | MiddlewareRef) -> MiddlewareBaseConfig:
+        """Get the configuration for middleware.
+
+        Args:
+            middleware_name: The name or reference of the middleware
+
+        Returns:
+            The configuration for the middleware
+
+        Raises:
+            ValueError: If the middleware is not found
+        """
+        if middleware_name not in self._middleware:
+            raise ValueError(f"Middleware `{middleware_name}` not found")
+
+        return self._middleware[middleware_name].config
+
     @override
     def get_user_manager(self):
         return UserManagerHolder(context=Context(self._context_state))
@@ -1108,6 +1218,10 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
                 elif component_instance.component_group == ComponentGroup.RETRIEVERS:
                     await self.add_retriever(component_instance.name,
                                              cast(RetrieverBaseConfig, component_instance.config))
+                # Instantiate middleware
+                elif component_instance.component_group == ComponentGroup.MIDDLEWARE:
+                    await self.add_middleware(component_instance.name,
+                                              cast(MiddlewareBaseConfig, component_instance.config))
                 # Instantiate a function group
                 elif component_instance.component_group == ComponentGroup.FUNCTION_GROUPS:
                     await self.add_function_group(component_instance.name,
@@ -1363,3 +1477,18 @@ class ChildBuilder(Builder):
     @override
     def get_function_group_dependencies(self, fn_name: str) -> FunctionDependencies:
         return self._workflow_builder.get_function_group_dependencies(fn_name)
+
+    @override
+    async def add_middleware(self, name: str | MiddlewareRef, config: MiddlewareBaseConfig) -> Middleware:
+        """Add middleware to the builder."""
+        return await self._workflow_builder.add_middleware(name, config)
+
+    @override
+    async def get_middleware(self, middleware_name: str | MiddlewareRef) -> Middleware:
+        """Get built middleware by name."""
+        return await self._workflow_builder.get_middleware(middleware_name)
+
+    @override
+    def get_middleware_config(self, middleware_name: str | MiddlewareRef) -> MiddlewareBaseConfig:
+        """Get the configuration for middleware."""
+        return self._workflow_builder.get_middleware_config(middleware_name)

nat/cli/register_workflow.py

@@ -38,6 +38,8 @@ from nat.cli.type_registry import LoggingMethodConfigT
 from nat.cli.type_registry import LoggingMethodRegisteredCallableT
 from nat.cli.type_registry import MemoryBuildCallableT
 from nat.cli.type_registry import MemoryRegisteredCallableT
+from nat.cli.type_registry import MiddlewareBuildCallableT
+from nat.cli.type_registry import MiddlewareRegisteredCallableT
 from nat.cli.type_registry import ObjectStoreBuildCallableT
 from nat.cli.type_registry import ObjectStoreRegisteredCallableT
 from nat.cli.type_registry import RegisteredLoggingMethod
@@ -65,6 +67,7 @@ from nat.data_models.function import FunctionConfigT
 from nat.data_models.function import FunctionGroupConfigT
 from nat.data_models.llm import LLMBaseConfigT
 from nat.data_models.memory import MemoryBaseConfigT
+from nat.data_models.middleware import MiddlewareBaseConfigT
 from nat.data_models.object_store import ObjectStoreBaseConfigT
 from nat.data_models.registry_handler import RegistryHandlerBaseConfigT
 from nat.data_models.retriever import RetrieverBaseConfigT
@@ -149,6 +152,10 @@ def register_function(config_type: type[FunctionConfigT],
                       framework_wrappers: list[LLMFrameworkEnum | str] | None = None):
     """
     Register a workflow with optional framework_wrappers for automatic profiler hooking.
+
+    Args:
+        config_type: The function configuration type
+        framework_wrappers: Optional list of framework wrappers for automatic profiler hooking
     """
 
     def register_function_inner(
@@ -211,6 +218,49 @@ def register_function_group(config_type: type[FunctionGroupConfigT],
     return register_function_group_inner
 
 
+def register_middleware(config_type: type[MiddlewareBaseConfigT]):
+    """
+    Register a middleware component.
+
+    Middleware provides middleware-style wrapping of calls with
+    preprocessing and postprocessing logic. They are built as components that can
+    be configured in YAML and referenced by name in configurations.
+
+    Args:
+        config_type: The middleware configuration type to register
+
+    Returns:
+        A decorator that wraps the build function as an async context manager
+    """
+
+    def register_middleware_inner(
+            fn: MiddlewareBuildCallableT[MiddlewareBaseConfigT]
+    ) -> MiddlewareRegisteredCallableT[MiddlewareBaseConfigT]:
+        from .type_registry import GlobalTypeRegistry
+        from .type_registry import RegisteredMiddlewareInfo
+
+        context_manager_fn = asynccontextmanager(fn)
+
+        discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
+                                                                component_type=ComponentEnum.MIDDLEWARE)
+
+        GlobalTypeRegistry.get().register_middleware(
+            RegisteredMiddlewareInfo(
+                full_type=config_type.full_type,
+                config_type=config_type,
+                build_fn=context_manager_fn,
+                discovery_metadata=discovery_metadata,
+            ))
+
+        return context_manager_fn
+
+    return register_middleware_inner
+
+
+# Compatibility alias for backwards compatibility
+register_function_middleware = register_middleware
+
+
 def register_llm_provider(config_type: type[LLMBaseConfigT]):
 
     def register_llm_provider_inner(