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

nat/agent/prompt_optimizer/register.py

@@ -51,7 +51,7 @@ async def prompt_optimizer_function(config: PromptOptimizerConfig, builder: Buil
         from .prompt import mutator_prompt
     except ImportError as exc:
         raise ImportError("langchain-core is not installed. Please install it to use MultiLLMPlanner.\n"
-                          "This error can be resolve by installing nvidia-nat[langchain]") from exc
+                          "This error can be resolve by installing \"nvidia-nat[langchain]\".") from exc
 
     llm = await builder.get_llm(config.optimizer_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
 
@@ -111,7 +111,7 @@ async def prompt_recombiner_function(config: PromptRecombinerConfig, builder: Bu
         from langchain_core.prompts import PromptTemplate
     except ImportError as exc:
         raise ImportError("langchain-core is not installed. Please install it to use MultiLLMPlanner.\n"
-                          "This error can be resolve by installing nvidia-nat[langchain].") from exc
+                          "This error can be resolve by installing \"nvidia-nat[langchain]\".") from exc
 
     llm = await builder.get_llm(config.optimizer_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
 

nat/agent/react_agent/register.py

@@ -25,6 +25,7 @@ from nat.cli.register_workflow import register_function
 from nat.data_models.agent import AgentBaseConfig
 from nat.data_models.api_server import ChatRequest
 from nat.data_models.api_server import ChatResponse
+from nat.data_models.api_server import Usage
 from nat.data_models.component_ref import FunctionGroupRef
 from nat.data_models.component_ref import FunctionRef
 from nat.data_models.optimizable import OptimizableField
@@ -149,7 +150,14 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
             # get and return the output from the state
             state = ReActGraphState(**state)
             output_message = state.messages[-1]
-            return ChatResponse.from_string(str(output_message.content))
+            content = str(output_message.content)
+
+            # Create usage statistics for the response
+            prompt_tokens = sum(len(str(msg.content).split()) for msg in input_message.messages)
+            completion_tokens = len(content.split()) if content else 0
+            total_tokens = prompt_tokens + completion_tokens
+            usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
+            return ChatResponse.from_string(content, usage=usage)
 
         except Exception as ex:
             logger.exception("%s ReAct Agent failed with exception: %s", AGENT_LOG_PREFIX, str(ex))

nat/agent/rewoo_agent/register.py

@@ -26,6 +26,7 @@ from nat.cli.register_workflow import register_function
 from nat.data_models.agent import AgentBaseConfig
 from nat.data_models.api_server import ChatRequest
 from nat.data_models.api_server import ChatResponse
+from nat.data_models.api_server import Usage
 from nat.data_models.component_ref import FunctionGroupRef
 from nat.data_models.component_ref import FunctionRef
 from nat.utils.type_converter import GlobalTypeConverter
@@ -157,7 +158,13 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
             # Ensure output_message is a string
             if isinstance(output_message, list | dict):
                 output_message = str(output_message)
-            return ChatResponse.from_string(output_message)
+
+            # Create usage statistics for the response
+            prompt_tokens = sum(len(str(msg.content).split()) for msg in input_message.messages)
+            completion_tokens = len(output_message.split()) if output_message else 0
+            total_tokens = prompt_tokens + completion_tokens
+            usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
+            return ChatResponse.from_string(output_message, usage=usage)
 
         except Exception as ex:
             logger.exception("ReWOO Agent failed with exception: %s", ex)

nat/authentication/oauth2/oauth2_auth_code_flow_provider.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import logging
+from collections.abc import Awaitable
 from collections.abc import Callable
 from datetime import UTC
 from datetime import datetime
@@ -35,10 +36,15 @@ logger = logging.getLogger(__name__)
 
 class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConfig]):
 
-    def __init__(self, config: OAuth2AuthCodeFlowProviderConfig):
+    def __init__(self, config: OAuth2AuthCodeFlowProviderConfig, token_storage=None):
         super().__init__(config)
-        self._authenticated_tokens: dict[str, AuthResult] = {}
         self._auth_callback = None
+        # Always use token storage - defaults to in-memory if not provided
+        if token_storage is None:
+            from nat.plugins.mcp.auth.token_storage import InMemoryTokenStorage
+            self._token_storage = InMemoryTokenStorage()
+        else:
+            self._token_storage = token_storage
 
     async def _attempt_token_refresh(self, user_id: str, auth_result: AuthResult) -> AuthResult | None:
         refresh_token = auth_result.raw.get("refresh_token")
@@ -61,7 +67,7 @@ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConf
                 raw=new_token_data,
             )
 
-            self._authenticated_tokens[user_id] = new_auth_result
+            await self._token_storage.store(user_id, new_auth_result)
         except httpx.HTTPStatusError:
             return None
         except httpx.RequestError:
@@ -74,26 +80,30 @@ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConf
 
     def _set_custom_auth_callback(self,
                                   auth_callback: Callable[[OAuth2AuthCodeFlowProviderConfig, AuthFlowType],
-                                                          AuthenticatedContext]):
+                                                          Awaitable[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)
+        context = Context.get()
+        if user_id is None and hasattr(context, "metadata") and hasattr(
+                context.metadata, "cookies") and context.metadata.cookies is not None:
+            session_id = context.metadata.cookies.get("nat-session", None)
             if not session_id:
                 raise RuntimeError("Authentication failed. No session ID found. Cannot identify user.")
 
             user_id = session_id
 
-        if user_id and user_id in self._authenticated_tokens:
-            auth_result = self._authenticated_tokens[user_id]
-            if not auth_result.is_expired():
-                return auth_result
+        if user_id:
+            # Try to retrieve from token storage
+            auth_result = await self._token_storage.retrieve(user_id)
+
+            if auth_result:
+                if not auth_result.is_expired():
+                    return auth_result
 
-            refreshed_auth_result = await self._attempt_token_refresh(user_id, auth_result)
-            if refreshed_auth_result:
-                return refreshed_auth_result
+                refreshed_auth_result = await self._attempt_token_refresh(user_id, auth_result)
+                if refreshed_auth_result:
+                    return refreshed_auth_result
 
         # Try getting callback from the context if that's not set, use the default callback
         try:
@@ -109,19 +119,22 @@ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConf
         except Exception as e:
             raise RuntimeError(f"Authentication callback failed: {e}") from e
 
-        auth_header = authenticated_context.headers.get("Authorization", "")
+        headers = authenticated_context.headers or {}
+        auth_header = headers.get("Authorization", "")
         if not auth_header.startswith("Bearer "):
             raise RuntimeError("Invalid Authorization header")
 
         token = auth_header.split(" ")[1]
 
+        # Safely access metadata
+        metadata = authenticated_context.metadata or {}
         auth_result = AuthResult(
             credentials=[BearerTokenCred(token=SecretStr(token))],
-            token_expires_at=authenticated_context.metadata.get("expires_at"),
-            raw=authenticated_context.metadata.get("raw_token"),
+            token_expires_at=metadata.get("expires_at"),
+            raw=metadata.get("raw_token") or {},
         )
 
         if user_id:
-            self._authenticated_tokens[user_id] = auth_result
+            await self._token_storage.store(user_id, auth_result)
 
         return auth_result

nat/builder/context.py

@@ -67,6 +67,8 @@ 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.workflow_run_id: ContextVar[str | None] = ContextVar("workflow_run_id", default=None)
+        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._metadata: ContextVar[RequestAttributes | None] = ContextVar("request_attributes", default=None)
@@ -120,14 +122,14 @@ class Context:
     @property
     def input_message(self):
         """
-            Retrieves the input message from the context state.
+        Retrieves the input message from the context state.
 
-            The input_message property is used to access the message stored in the
-            context state. This property returns the message as it is currently
-            maintained in the context.
+        The input_message property is used to access the message stored in the
+        context state. This property returns the message as it is currently
+        maintained in the context.
 
-            Returns:
-                str: The input message retrieved from the context state.
+        Returns:
+            str: The input message retrieved from the context state.
         """
         return self._context_state.input_message.get()
 
@@ -196,6 +198,20 @@ class Context:
         """
         return self._context_state.user_message_id.get()
 
+    @property
+    def workflow_run_id(self) -> str | None:
+        """
+        Returns a stable identifier for the current workflow/agent invocation (UUID string).
+        """
+        return self._context_state.workflow_run_id.get()
+
+    @property
+    def workflow_trace_id(self) -> int | None:
+        """
+        Returns the 128-bit trace identifier for the current run, used as the OpenTelemetry trace_id.
+        """
+        return self._context_state.workflow_trace_id.get()
+
     @contextmanager
     def push_active_function(self,
                              function_name: str,

nat/cli/commands/mcp/mcp.py

@@ -194,7 +194,7 @@ async def _create_mcp_client_config(
     auth_user_id: str | None,
     auth_scopes: list[str] | None,
 ):
-    from nat.plugins.mcp.client_impl import MCPClientConfig
+    from nat.plugins.mcp.client_config import MCPClientConfig
 
     if url and transport == "streamable-http" and auth_redirect_uri:
         try:
@@ -236,8 +236,8 @@ async def list_tools_via_function_group(
     try:
         # Ensure the registration side-effects are loaded
         from nat.builder.workflow_builder import WorkflowBuilder
-        from nat.plugins.mcp.client_impl import MCPClientConfig
-        from nat.plugins.mcp.client_impl import MCPServerConfig
+        from nat.plugins.mcp.client_config import MCPClientConfig
+        from nat.plugins.mcp.client_config import MCPServerConfig
     except ImportError:
         click.echo(
             "MCP client functionality requires nvidia-nat-mcp package. Install with: uv pip install nvidia-nat-mcp",
@@ -297,7 +297,7 @@ async def list_tools_via_function_group(
             if fn is not None:
                 tools.append(to_tool_entry(full, fn))
         else:
-            for full, fn in (await fns).items():
+            for full, fn in fns.items():
                 tools.append(to_tool_entry(full, fn))
 
     return tools
@@ -826,8 +826,8 @@ async def call_tool_and_print(command: str | None,
 
     try:
         from nat.builder.workflow_builder import WorkflowBuilder
-        from nat.plugins.mcp.client_impl import MCPClientConfig
-        from nat.plugins.mcp.client_impl import MCPServerConfig
+        from nat.plugins.mcp.client_config import MCPClientConfig
+        from nat.plugins.mcp.client_config import MCPServerConfig
     except ImportError:
         click.echo(
             "MCP client functionality requires nvidia-nat-mcp package. Install with: uv pip install nvidia-nat-mcp",

nat/cli/commands/workflow/templates/config.yml.j2

@@ -1,15 +1,17 @@
-general:
-  logging:
-    console:
-      _type: console
-      level: WARN
+functions:
+  current_datetime:
+    _type: current_datetime
+  {{python_safe_workflow_name}}:
+    _type: {{python_safe_workflow_name}}
+    prefix: "Hello:"
 
-  front_end:
-    _type: fastapi
-
-  front_end:
-    _type: console
+llms:
+  nim_llm:
+    _type: nim
+    model_name: meta/llama-3.1-70b-instruct
+    temperature: 0.0
 
 workflow:
-  _type: {{workflow_name}}
-  parameter: default_value
+  _type: react_agent
+  llm_name: nim_llm
+  tool_names: [current_datetime, {{python_safe_workflow_name}}]

nat/cli/commands/workflow/templates/register.py.j2

@@ -1,4 +1,4 @@
 # flake8: noqa
 
-# Import any tools which need to be automatically registered here
-from {{package_name}} import {{workflow_name}}_function
+# Import the generated workflow function to trigger registration
+from .{{package_name}} import {{ python_safe_workflow_name }}_function

nat/cli/commands/workflow/templates/workflow.py.j2

@@ -3,6 +3,7 @@ import logging
 from pydantic import Field
 
 from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.builder.function_info import FunctionInfo
 from nat.cli.register_workflow import register_function
 from nat.data_models.function import FunctionBaseConfig
@@ -12,25 +13,38 @@ logger = logging.getLogger(__name__)
 
 class {{ workflow_class_name }}(FunctionBaseConfig, name="{{ workflow_name }}"):
     """
-    {{workflow_description}}
+    {{ workflow_description }}
     """
-    # Add your custom configuration parameters here
-    parameter: str = Field(default="default_value", description="Notional description for this parameter")
-
-
-@register_function(config_type={{ workflow_class_name }})
-async def {{ python_safe_workflow_name }}_function(
-    config: {{ workflow_class_name }}, builder: Builder
-):
-    # Implement your function logic here
-    async def _response_fn(input_message: str) -> str:
-        # Process the input_message and generate output
-        output_message = f"Hello from {{ workflow_name }} workflow! You said: {input_message}"
-        return output_message
-
-    try:
-        yield FunctionInfo.create(single_fn=_response_fn)
-    except GeneratorExit:
-        logger.warning("Function exited early!")
-    finally:
-        logger.info("Cleaning up {{ workflow_name }} workflow.")
+    prefix: str = Field(default="Echo:", description="Prefix to add before the echoed text.")
+
+
+@register_function(config_type={{ workflow_class_name }}, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def {{ python_safe_workflow_name }}_function(config: {{ workflow_class_name }}, builder: Builder):
+    """
+    Registers a function (addressable via `{{ workflow_name }}` in the configuration).
+    This registration ensures a static mapping of the function type, `{{ workflow_name }}`, to the `{{ workflow_class_name }}` configuration object.
+
+    Args:
+        config ({{ workflow_class_name }}): The configuration for the function.
+        builder (Builder): The builder object.
+
+    Returns:
+        FunctionInfo: The function info object for the function.
+    """
+
+    # Define the function that will be registered.
+    async def _echo(text: str) -> str:
+        """
+        Takes a text input and echoes back with a pre-defined prefix.
+
+        Args:
+            text (str): The text to echo back.
+
+        Returns:
+            str: The text with the prefix.
+        """
+        return f"{config.prefix} {text}"
+
+    # The callable is wrapped in a FunctionInfo object.
+    # The description parameter is used to describe the function.
+    yield FunctionInfo.from_fn(_echo, description=_echo.__doc__)

nat/cli/commands/workflow/workflow_commands.py

@@ -27,6 +27,50 @@ from jinja2 import FileSystemLoader
 logger = logging.getLogger(__name__)
 
 
+def _get_nat_version() -> str | None:
+    """
+    Get the current NAT version.
+
+    Returns:
+        str: The NAT version intended for use in a dependency string.
+        None: If the NAT version is not found.
+    """
+    from nat.cli.entrypoint import get_version
+
+    current_version = get_version()
+    if current_version == "unknown":
+        return None
+
+    version_parts = current_version.split(".")
+    if len(version_parts) < 3:
+        # If the version somehow doesn't have three parts, return the full version
+        return current_version
+
+    patch = version_parts[2]
+    try:
+        # If the patch is a number, keep only the major and minor parts
+        # Useful for stable releases and adheres to semantic versioning
+        _ = int(patch)
+        digits_to_keep = 2
+    except ValueError:
+        # If the patch is not a number, keep all three digits
+        # Useful for pre-release versions (and nightly builds)
+        digits_to_keep = 3
+
+    return ".".join(version_parts[:digits_to_keep])
+
+
+def _is_nat_version_prerelease() -> bool:
+    """
+    Check if the NAT version is a prerelease.
+    """
+    version = _get_nat_version()
+    if version is None:
+        return False
+
+    return len(version.split(".")) >= 3
+
+
 def _get_nat_dependency(versioned: bool = True) -> str:
     """
     Get the NAT dependency string with version.
@@ -44,16 +88,12 @@ def _get_nat_dependency(versioned: bool = True) -> str:
         logger.debug("Using unversioned NAT dependency: %s", dependency)
         return dependency
 
-    # Get the current NAT version
-    from nat.cli.entrypoint import get_version
-    current_version = get_version()
-    if current_version == "unknown":
-        logger.warning("Could not detect NAT version, using unversioned dependency")
+    version = _get_nat_version()
+    if version is None:
+        logger.debug("Could not detect NAT version, using unversioned dependency: %s", dependency)
         return dependency
 
-    # Extract major.minor (e.g., "1.2.3" -> "1.2")
-    major_minor = ".".join(current_version.split(".")[:2])
-    dependency += f"~={major_minor}"
+    dependency += f"~={version}"
     logger.debug("Using NAT dependency: %s", dependency)
     return dependency
 
@@ -219,12 +259,16 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
             install_cmd = ['uv', 'pip', 'install', '-e', str(new_workflow_dir)]
         else:
             install_cmd = ['pip', 'install', '-e', str(new_workflow_dir)]
+            if _is_nat_version_prerelease():
+                install_cmd.insert(2, "--pre")
+
+        python_safe_workflow_name = workflow_name.replace("-", "_")
 
         # List of templates and their destinations
         files_to_render = {
             'pyproject.toml.j2': new_workflow_dir / 'pyproject.toml',
             'register.py.j2': base_dir / 'register.py',
-            'workflow.py.j2': base_dir / f'{workflow_name}_function.py',
+            'workflow.py.j2': base_dir / f'{python_safe_workflow_name}.py',
             '__init__.py.j2': base_dir / '__init__.py',
             'config.yml.j2': configs_dir / 'config.yml',
         }
@@ -233,7 +277,7 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
         context = {
             'editable': editable,
             'workflow_name': workflow_name,
-            'python_safe_workflow_name': workflow_name.replace("-", "_"),
+            'python_safe_workflow_name': python_safe_workflow_name,
             'package_name': package_name,
             'rel_path_to_repo_root': rel_path_to_repo_root,
             'workflow_class_name': f"{_generate_valid_classname(workflow_name)}FunctionConfig",

nat/cli/main.py

@@ -30,6 +30,9 @@ def run_cli():
     import os
     import sys
 
+    # Suppress warnings from transformers
+    os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+
     parent_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
 
     if (parent_dir not in sys.path):