nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

nat/authentication/oauth2/oauth2_auth_code_flow_provider.py

@@ -13,56 +13,71 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+from collections.abc import Callable
+from datetime import UTC
 from datetime import datetime
-from datetime import timezone
 
+import httpx
 from authlib.integrations.httpx_client import OAuth2Client as AuthlibOAuth2Client
 from pydantic import SecretStr
 
 from nat.authentication.interfaces import AuthProviderBase
 from nat.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
 from nat.builder.context import Context
+from nat.data_models.authentication import AuthenticatedContext
 from nat.data_models.authentication import AuthFlowType
 from nat.data_models.authentication import AuthResult
 from nat.data_models.authentication import BearerTokenCred
 
+logger = logging.getLogger(__name__)
+
 
 class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConfig]):
 
     def __init__(self, config: OAuth2AuthCodeFlowProviderConfig):
         super().__init__(config)
         self._authenticated_tokens: dict[str, AuthResult] = {}
-        self._context = Context.get()
+        self._auth_callback = None
 
     async def _attempt_token_refresh(self, user_id: str, auth_result: AuthResult) -> AuthResult | None:
         refresh_token = auth_result.raw.get("refresh_token")
         if not isinstance(refresh_token, str):
             return None
 
-        with AuthlibOAuth2Client(
-                client_id=self.config.client_id,
-                client_secret=self.config.client_secret,
-        ) as client:
-            try:
+        try:
+            with AuthlibOAuth2Client(
+                    client_id=self.config.client_id,
+                    client_secret=self.config.client_secret,
+            ) as client:
                 new_token_data = client.refresh_token(self.config.token_url, refresh_token=refresh_token)
-            except Exception:
-                # On any failure, we'll fall back to the full auth flow.
-                return None
 
-        expires_at_ts = new_token_data.get("expires_at")
-        new_expires_at = datetime.fromtimestamp(expires_at_ts, tz=timezone.utc) if expires_at_ts else None
+                expires_at_ts = new_token_data.get("expires_at")
+                new_expires_at = datetime.fromtimestamp(expires_at_ts, tz=UTC) if expires_at_ts else None
 
-        new_auth_result = AuthResult(
-            credentials=[BearerTokenCred(token=SecretStr(new_token_data["access_token"]))],
-            token_expires_at=new_expires_at,
-            raw=new_token_data,
-        )
+            new_auth_result = AuthResult(
+                credentials=[BearerTokenCred(token=SecretStr(new_token_data["access_token"]))],
+                token_expires_at=new_expires_at,
+                raw=new_token_data,
+            )
 
-        self._authenticated_tokens[user_id] = new_auth_result
+            self._authenticated_tokens[user_id] = new_auth_result
+        except httpx.HTTPStatusError:
+            return None
+        except httpx.RequestError:
+            return None
+        except Exception:
+            # On any other failure, we'll fall back to the full auth flow.
+            return None
 
         return new_auth_result
 
-    async def authenticate(self, user_id: str | None = None) -> AuthResult:
+    def _set_custom_auth_callback(self,
+                                  auth_callback: Callable[[OAuth2AuthCodeFlowProviderConfig, AuthFlowType],
+                                                          AuthenticatedContext]):
+        self._auth_callback = auth_callback
+
+    async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult:
         if user_id is None and hasattr(Context.get(), "metadata") and hasattr(
                 Context.get().metadata, "cookies") and Context.get().metadata.cookies is not None:
             session_id = Context.get().metadata.cookies.get("nat-session", None)
@@ -80,7 +95,12 @@ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConf
             if refreshed_auth_result:
                 return refreshed_auth_result
 
-        auth_callback = self._context.user_auth_callback
+        # Try getting callback from the context if that's not set, use the default callback
+        try:
+            auth_callback = Context.get().user_auth_callback
+        except RuntimeError:
+            auth_callback = self._auth_callback
+
         if not auth_callback:
             raise RuntimeError("Authentication callback not set on Context.")
 

nat/authentication/oauth2/oauth2_resource_server_config.py

@@ -0,0 +1,124 @@
+# 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 urllib.parse import urlparse
+
+from pydantic import Field
+from pydantic import field_validator
+from pydantic import model_validator
+
+from nat.data_models.authentication import AuthProviderBaseConfig
+
+
+class OAuth2ResourceServerConfig(AuthProviderBaseConfig, name="oauth2_resource_server"):
+    """OAuth 2.0 Resource Server authentication configuration.
+
+    Supports:
+      • JWT access tokens via JWKS / OIDC Discovery / issuer fallback
+      • Opaque access tokens via RFC 7662 introspection
+    """
+
+    issuer_url: str = Field(
+        description=("The unique issuer identifier for an authorization server. "
+                     "Required for validation and used to derive the default JWKS URI "
+                     "(<issuer_url>/.well-known/jwks.json) if `jwks_uri` and `discovery_url` are not provided."), )
+    scopes: list[str] = Field(
+        default_factory=list,
+        description="Scopes required by this API. Validation ensures the token grants all listed scopes.",
+    )
+    audience: str | None = Field(
+        default=None,
+        description=(
+            "Expected audience (`aud`) claim for this API. If set, validation will reject tokens without this audience."
+        ),
+    )
+
+    # JWT verification params
+    jwks_uri: str | None = Field(
+        default=None,
+        description=("Direct JWKS endpoint URI for JWT signature verification. "
+                     "Optional if discovery or issuer is provided."),
+    )
+    discovery_url: str | None = Field(
+        default=None,
+        description=("OIDC discovery metadata URL. Used to automatically resolve JWKS and introspection endpoints."),
+    )
+
+    # Opaque token (introspection) params
+    introspection_endpoint: str | None = Field(
+        default=None,
+        description=("RFC 7662 token introspection endpoint. "
+                     "Required for opaque token validation and must be used with `client_id` and `client_secret`."),
+    )
+    client_id: str | None = Field(
+        default=None,
+        description="OAuth2 client ID for authenticating to the introspection endpoint (opaque token validation).",
+    )
+    client_secret: str | None = Field(
+        default=None,
+        description="OAuth2 client secret for authenticating to the introspection endpoint (opaque token validation).",
+    )
+
+    @staticmethod
+    def _is_https_or_localhost(url: str) -> bool:
+        try:
+            value = urlparse(url)
+            if not value.scheme or not value.netloc:
+                return False
+            if value.scheme == "https":
+                return True
+            return value.scheme == "http" and (value.hostname in {"localhost", "127.0.0.1", "::1"})
+        except Exception:
+            return False
+
+    @field_validator("issuer_url", "jwks_uri", "discovery_url", "introspection_endpoint")
+    @classmethod
+    def _require_valid_url(cls, value: str | None, info):
+        if value is None:
+            return value
+        if not cls._is_https_or_localhost(value):
+            raise ValueError(f"{info.field_name} must be HTTPS (http allowed only for localhost). Got: {value}")
+        return value
+
+    # ---------- Cross-field validation: ensure at least one viable path ----------
+
+    @model_validator(mode="after")
+    def _ensure_verification_path(self):
+        """
+        JWT path viable if any of: jwks_uri OR discovery_url OR issuer_url (fallback JWKS).
+        Opaque path viable if: introspection_endpoint AND client_id AND client_secret.
+        """
+        has_jwt_path = bool(self.jwks_uri or self.discovery_url or self.issuer_url)
+        has_opaque_path = bool(self.introspection_endpoint and self.client_id and self.client_secret)
+
+        # If introspection endpoint is set, enforce creds are present
+        if self.introspection_endpoint:
+            missing = []
+            if not self.client_id:
+                missing.append("client_id")
+            if not self.client_secret:
+                missing.append("client_secret")
+            if missing:
+                raise ValueError(
+                    f"introspection_endpoint configured but missing required credentials: {', '.join(missing)}")
+
+        # Require at least one path
+        if not (has_jwt_path or has_opaque_path):
+            raise ValueError("Invalid configuration: no verification method available. "
+                             "Configure one of the following:\n"
+                             "  • JWT path: set jwks_uri OR discovery_url OR issuer_url (for JWKS fallback)\n"
+                             "  • Opaque path: set introspection_endpoint + client_id + client_secret")
+
+        return self

nat/authentication/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
 
 from nat.authentication.api_key import register as register_api_key

nat/builder/builder.py

@@ -24,9 +24,11 @@ from nat.authentication.interfaces import AuthProviderBase
 from nat.builder.context import Context
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.builder.function import Function
+from nat.builder.function import FunctionGroup
 from nat.data_models.authentication import AuthProviderBaseConfig
 from nat.data_models.component_ref import AuthenticationRef
 from nat.data_models.component_ref import EmbedderRef
+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
@@ -36,20 +38,25 @@ from nat.data_models.component_ref import TTCStrategyRef
 from nat.data_models.embedder import EmbedderBaseConfig
 from nat.data_models.evaluator import EvaluatorBaseConfig
 from nat.data_models.function import FunctionBaseConfig
+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.object_store import ObjectStoreBaseConfig
 from nat.data_models.retriever import RetrieverBaseConfig
 from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.decorators.experimental_warning_decorator import experimental
 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.object_store.interfaces import ObjectStore
 from nat.retriever.interface import Retriever
 
+if typing.TYPE_CHECKING:
+    from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
 
-class UserManagerHolder():
+
+class UserManagerHolder:
 
     def __init__(self, context: Context) -> None:
         self._context = context
@@ -58,24 +65,40 @@ class UserManagerHolder():
         return self._context.user_manager.get_id()
 
 
-class Builder(ABC):  # pylint: disable=too-many-public-methods
+class Builder(ABC):
 
     @abstractmethod
     async def add_function(self, name: str | FunctionRef, config: FunctionBaseConfig) -> Function:
         pass
 
     @abstractmethod
-    def get_function(self, name: str | FunctionRef) -> Function:
+    async def add_function_group(self, name: str | FunctionGroupRef, config: FunctionGroupBaseConfig) -> FunctionGroup:
+        pass
+
+    @abstractmethod
+    async def get_function(self, name: str | FunctionRef) -> Function:
         pass
 
-    def get_functions(self, function_names: Sequence[str | FunctionRef]) -> list[Function]:
+    @abstractmethod
+    async def get_function_group(self, name: str | FunctionGroupRef) -> FunctionGroup:
+        pass
 
-        return [self.get_function(name) for name in function_names]
+    async def get_functions(self, function_names: Sequence[str | FunctionRef]) -> list[Function]:
+        tasks = [self.get_function(name) for name in function_names]
+        return list(await asyncio.gather(*tasks, return_exceptions=False))
+
+    async def get_function_groups(self, function_group_names: Sequence[str | FunctionGroupRef]) -> list[FunctionGroup]:
+        tasks = [self.get_function_group(name) for name in function_group_names]
+        return list(await asyncio.gather(*tasks, return_exceptions=False))
 
     @abstractmethod
     def get_function_config(self, name: str | FunctionRef) -> FunctionBaseConfig:
         pass
 
+    @abstractmethod
+    def get_function_group_config(self, name: str | FunctionGroupRef) -> FunctionGroupBaseConfig:
+        pass
+
     @abstractmethod
     async def set_workflow(self, config: FunctionBaseConfig) -> Function:
         pass
@@ -88,17 +111,18 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
     def get_workflow_config(self) -> FunctionBaseConfig:
         pass
 
-    def get_tools(self, tool_names: Sequence[str | FunctionRef],
-                  wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
-
-        return [self.get_tool(fn_name=n, wrapper_type=wrapper_type) for n in tool_names]
+    @abstractmethod
+    async def get_tools(self,
+                        tool_names: Sequence[str | FunctionRef | FunctionGroupRef],
+                        wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
+        pass
 
     @abstractmethod
-    def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
+    async def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
         pass
 
     @abstractmethod
-    async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig):
+    async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig) -> typing.Any:
         pass
 
     @abstractmethod
@@ -119,7 +143,9 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
         pass
 
     @abstractmethod
-    async def add_auth_provider(self, name: str | AuthenticationRef, config: AuthProviderBaseConfig):
+    @experimental(feature_name="Authentication")
+    async def add_auth_provider(self, name: str | AuthenticationRef,
+                                config: AuthProviderBaseConfig) -> AuthProviderBase:
         pass
 
     @abstractmethod
@@ -135,7 +161,7 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
         return list(auth_providers)
 
     @abstractmethod
-    async def add_object_store(self, name: str | ObjectStoreRef, config: ObjectStoreBaseConfig):
+    async def add_object_store(self, name: str | ObjectStoreRef, config: ObjectStoreBaseConfig) -> ObjectStore:
         pass
 
     async def get_object_store_clients(self, object_store_names: Sequence[str | ObjectStoreRef]) -> list[ObjectStore]:
@@ -153,7 +179,7 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
         pass
 
     @abstractmethod
-    async def add_embedder(self, name: str | EmbedderRef, config: EmbedderBaseConfig):
+    async def add_embedder(self, name: str | EmbedderRef, config: EmbedderBaseConfig) -> None:
         pass
 
     async def get_embedders(self, embedder_names: Sequence[str | EmbedderRef],
@@ -174,17 +200,18 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
         pass
 
     @abstractmethod
-    async def add_memory_client(self, name: str | MemoryRef, config: MemoryBaseConfig):
+    async def add_memory_client(self, name: str | MemoryRef, config: MemoryBaseConfig) -> MemoryEditor:
         pass
 
-    def get_memory_clients(self, memory_names: Sequence[str | MemoryRef]) -> list[MemoryEditor]:
+    async def get_memory_clients(self, memory_names: Sequence[str | MemoryRef]) -> list[MemoryEditor]:
         """
         Return a list of memory clients for the specified names.
         """
-        return [self.get_memory_client(n) for n in memory_names]
+        tasks = [self.get_memory_client(n) for n in memory_names]
+        return list(await asyncio.gather(*tasks, return_exceptions=False))
 
     @abstractmethod
-    def get_memory_client(self, memory_name: str | MemoryRef) -> MemoryEditor:
+    async def get_memory_client(self, memory_name: str | MemoryRef) -> MemoryEditor:
         """
         Return the instantiated memory client for the given name.
         """
@@ -195,12 +222,12 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
         pass
 
     @abstractmethod
-    async def add_retriever(self, name: str | RetrieverRef, config: RetrieverBaseConfig):
+    async def add_retriever(self, name: str | RetrieverRef, config: RetrieverBaseConfig) -> None:
         pass
 
     async def get_retrievers(self,
                              retriever_names: Sequence[str | RetrieverRef],
-                             wrapper_type: LLMFrameworkEnum | str | None = None):
+                             wrapper_type: LLMFrameworkEnum | str | None = None) -> list[Retriever]:
 
         tasks = [self.get_retriever(n, wrapper_type=wrapper_type) for n in retriever_names]
 
@@ -232,14 +259,15 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
         pass
 
     @abstractmethod
-    async def add_ttc_strategy(self, name: str | str, config: TTCStrategyBaseConfig):
+    @experimental(feature_name="TTC")
+    async def add_ttc_strategy(self, name: str | TTCStrategyRef, config: TTCStrategyBaseConfig):
         pass
 
     @abstractmethod
     async def get_ttc_strategy(self,
                                strategy_name: str | TTCStrategyRef,
                                pipeline_type: PipelineTypeEnum,
-                               stage_type: StageTypeEnum):
+                               stage_type: StageTypeEnum) -> "StrategyBase":
         pass
 
     @abstractmethod
@@ -257,8 +285,12 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
     def get_function_dependencies(self, fn_name: str) -> FunctionDependencies:
         pass
 
+    @abstractmethod
+    def get_function_group_dependencies(self, fn_name: str) -> FunctionDependencies:
+        pass
+
 
-class EvalBuilder(Builder):
+class EvalBuilder(ABC):
 
     @abstractmethod
     async def add_evaluator(self, name: str, config: EvaluatorBaseConfig):
@@ -281,5 +313,5 @@ class EvalBuilder(Builder):
         pass
 
     @abstractmethod
-    def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
+    async def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
         pass

nat/builder/component_utils.py

@@ -30,6 +30,7 @@ from nat.data_models.component_ref import generate_instance_id
 from nat.data_models.config import Config
 from nat.data_models.embedder import EmbedderBaseConfig
 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.object_store import ObjectStoreBaseConfig
@@ -48,6 +49,7 @@ _component_group_order = [
     ComponentGroup.OBJECT_STORES,
     ComponentGroup.RETRIEVERS,
     ComponentGroup.TTC_STRATEGIES,
+    ComponentGroup.FUNCTION_GROUPS,
     ComponentGroup.FUNCTIONS,
 ]
 
@@ -107,6 +109,8 @@ def group_from_component(component: TypedBaseModel) -> ComponentGroup | None:
         return ComponentGroup.EMBEDDERS
     if (isinstance(component, FunctionBaseConfig)):
         return ComponentGroup.FUNCTIONS
+    if (isinstance(component, FunctionGroupBaseConfig)):
+        return ComponentGroup.FUNCTION_GROUPS
     if (isinstance(component, LLMBaseConfig)):
         return ComponentGroup.LLMS
     if (isinstance(component, MemoryBaseConfig)):
@@ -154,7 +158,7 @@ def recursive_componentref_discovery(cls: TypedBaseModel, value: typing.Any,
             yield from recursive_componentref_discovery(cls, field_data, field_info.annotation)
     if (decomposed_type.is_union):
         for arg in decomposed_type.args:
-            if arg is typing.Any or (isinstance(value, DecomposedType(arg).root)):
+            if arg is typing.Any or DecomposedType(arg).is_instance(value):
                 yield from recursive_componentref_discovery(cls, value, arg)
     else:
         for arg in decomposed_type.args:
@@ -174,7 +178,7 @@ def update_dependency_graph(config: "Config", instance_config: TypedBaseModel,
         nx.DiGraph: An dependency graph that has been updated with the provided runtime instance.
     """
 
-    for field_name, field_info in instance_config.model_fields.items():
+    for field_name, field_info in type(instance_config).model_fields.items():
 
         for instance_id, value_node in recursive_componentref_discovery(
                 instance_config,
@@ -254,9 +258,9 @@ def build_dependency_sequence(config: "Config") -> list[ComponentInstanceData]:
             runtime instance references.
     """
 
-    total_node_count = len(config.embedders) + len(config.functions) + 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
+    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
 
     dependency_map: dict
     dependency_graph: nx.DiGraph

nat/builder/context.py

@@ -31,6 +31,7 @@ from nat.data_models.intermediate_step import IntermediateStep
 from nat.data_models.intermediate_step import IntermediateStepPayload
 from nat.data_models.intermediate_step import IntermediateStepType
 from nat.data_models.intermediate_step import StreamEventData
+from nat.data_models.intermediate_step import TraceMetadata
 from nat.data_models.invocation_node import InvocationNode
 from nat.runtime.user_metadata import RequestAttributes
 from nat.utils.reactive.subject import Subject
@@ -38,13 +39,13 @@ from nat.utils.reactive.subject import Subject
 
 class Singleton(type):
 
-    def __init__(cls, name, bases, dict):  # pylint: disable=W0622
-        super(Singleton, cls).__init__(name, bases, dict)
+    def __init__(cls, name, bases, dict):
+        super().__init__(name, bases, dict)
         cls.instance = None
 
     def __call__(cls, *args, **kw):
         if cls.instance is None:
-            cls.instance = super(Singleton, cls).__call__(*args, **kw)
+            cls.instance = super().__call__(*args, **kw)
         return cls.instance
 
 
@@ -65,14 +66,13 @@ class ContextState(metaclass=Singleton):
 
     def __init__(self):
         self.conversation_id: ContextVar[str | None] = ContextVar("conversation_id", default=None)
+        self.user_message_id: ContextVar[str | None] = ContextVar("user_message_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.metadata: ContextVar[RequestAttributes] = ContextVar("request_attributes", default=RequestAttributes())
-        self.event_stream: ContextVar[Subject[IntermediateStep] | None] = ContextVar("event_stream", default=Subject())
-        self.active_function: ContextVar[InvocationNode] = ContextVar("active_function",
-                                                                      default=InvocationNode(function_id="root",
-                                                                                             function_name="root"))
-        self.active_span_id_stack: ContextVar[list[str]] = ContextVar("active_span_id_stack", default=["root"])
+        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)
+        self._active_span_id_stack: ContextVar[list[str] | None] = ContextVar("active_span_id_stack", default=None)
 
         # Default is a lambda no-op which returns NoneType
         self.user_input_callback: ContextVar[Callable[[InteractionPrompt], Awaitable[HumanResponse | None]]
@@ -83,6 +83,30 @@ class ContextState(metaclass=Singleton):
                                                      Awaitable[AuthenticatedContext]]
                                             | None] = ContextVar("user_auth_callback", default=None)
 
+    @property
+    def metadata(self) -> ContextVar[RequestAttributes]:
+        if self._metadata.get() is None:
+            self._metadata.set(RequestAttributes())
+        return typing.cast(ContextVar[RequestAttributes], self._metadata)
+
+    @property
+    def active_function(self) -> ContextVar[InvocationNode]:
+        if self._active_function.get() is None:
+            self._active_function.set(InvocationNode(function_id="root", function_name="root"))
+        return typing.cast(ContextVar[InvocationNode], self._active_function)
+
+    @property
+    def event_stream(self) -> ContextVar[Subject[IntermediateStep]]:
+        if self._event_stream.get() is None:
+            self._event_stream.set(Subject())
+        return typing.cast(ContextVar[Subject[IntermediateStep]], self._event_stream)
+
+    @property
+    def active_span_id_stack(self) -> ContextVar[list[str]]:
+        if self._active_span_id_stack.get() is None:
+            self._active_span_id_stack.set(["root"])
+        return typing.cast(ContextVar[list[str]], self._active_span_id_stack)
+
     @staticmethod
     def get() -> "ContextState":
         return ContextState()
@@ -165,8 +189,18 @@ class Context:
         """
         return self._context_state.conversation_id.get()
 
+    @property
+    def user_message_id(self) -> str | None:
+        """
+        This property retrieves the user message ID which is the unique identifier for the current user message.
+        """
+        return self._context_state.user_message_id.get()
+
     @contextmanager
-    def push_active_function(self, function_name: str, input_data: typing.Any | None):
+    def push_active_function(self,
+                             function_name: str,
+                             input_data: typing.Any | None,
+                             metadata: dict[str, typing.Any] | TraceMetadata | None = None):
         """
         Set the 'active_function' in context, push an invocation node,
         AND create an OTel child span for that function call.
@@ -187,7 +221,8 @@ class Context:
             IntermediateStepPayload(UUID=current_function_id,
                                     event_type=IntermediateStepType.FUNCTION_START,
                                     name=function_name,
-                                    data=StreamEventData(input=input_data)))
+                                    data=StreamEventData(input=input_data),
+                                    metadata=metadata))
 
         manager = ActiveFunctionContextManager()
 

nat/builder/eval_builder.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import dataclasses
 import logging
 from contextlib import asynccontextmanager
@@ -61,7 +62,7 @@ class WorkflowEvalBuilder(WorkflowBuilder, EvalBuilder):
             # Store the evaluator
             self._evaluators[name] = ConfiguredEvaluator(config=config, instance=info_obj)
         except Exception as e:
-            logger.error("Error %s adding evaluator `%s` with config `%s`", e, name, config, exc_info=True)
+            logger.error("Error %s adding evaluator `%s` with config `%s`", e, name, config)
             raise
 
     @override
@@ -90,17 +91,20 @@ class WorkflowEvalBuilder(WorkflowBuilder, EvalBuilder):
         return self.eval_general_config.output_dir
 
     @override
-    def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str):
-        tools = []
+    async def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str):
         tool_wrapper_reg = self._registry.get_tool_wrapper(llm_framework=wrapper_type)
-        for fn_name in self._functions:
-            fn = self.get_function(fn_name)
+
+        async def get_tool(fn_name: str):
+            fn = await self.get_function(fn_name)
             try:
-                tools.append(tool_wrapper_reg.build_fn(fn_name, fn, self))
+                return tool_wrapper_reg.build_fn(fn_name, fn, self)
             except Exception:
-                logger.exception("Error fetching tool `%s`", fn_name, exc_info=True)
+                logger.exception("Error fetching tool `%s`", fn_name)
+                return None
 
-        return tools
+        tasks = [get_tool(fn_name) for fn_name in self._functions]
+        tools = await asyncio.gather(*tasks, return_exceptions=False)
+        return [tool for tool in tools if tool is not None]
 
     def _log_build_failure_evaluator(self,
                                      failing_evaluator_name: str,
@@ -127,11 +131,12 @@ class WorkflowEvalBuilder(WorkflowBuilder, EvalBuilder):
                                 remaining_components,
                                 original_error)
 
-    async def populate_builder(self, config: Config):
+    @override
+    async def populate_builder(self, config: Config, skip_workflow: bool = False):
         # Skip setting workflow if workflow config is EmptyFunctionConfig
-        skip_workflow = isinstance(config.workflow, EmptyFunctionConfig)
+        skip_workflow = skip_workflow or isinstance(config.workflow, EmptyFunctionConfig)
 
-        await super().populate_builder(config, skip_workflow)
+        await super().populate_builder(config, skip_workflow=skip_workflow)
 
         # Initialize progress tracking for evaluators
         completed_evaluators = []