nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/runtime/runner.py

@@ -0,0 +1,292 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import typing
+import uuid
+from enum import Enum
+
+from nat.builder.context import Context
+from nat.builder.context import ContextState
+from nat.builder.function import Function
+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.observability.exporter_manager import ExporterManager
+from nat.utils.reactive.subject import Subject
+
+logger = logging.getLogger(__name__)
+
+
+class UserManagerBase:
+    pass
+
+
+class RunnerState(Enum):
+    UNINITIALIZED = 0
+    INITIALIZED = 1
+    RUNNING = 2
+    COMPLETED = 3
+    FAILED = 4
+
+
+_T = typing.TypeVar("_T")
+
+
+class Runner:
+
+    def __init__(self,
+                 input_message: typing.Any,
+                 entry_fn: Function,
+                 context_state: ContextState,
+                 exporter_manager: ExporterManager):
+        """
+        The Runner class is used to run a workflow. It handles converting input and output data types and running the
+        workflow with the specified concurrency.
+
+        Parameters
+        ----------
+        input_message : typing.Any
+            The input message to the workflow
+        entry_fn : Function
+            The entry function to the workflow
+        context_state : ContextState
+            The context state to use
+        exporter_manager : ExporterManager
+            The exporter manager to use
+        """
+
+        if (entry_fn is None):
+            raise ValueError("entry_fn cannot be None")
+
+        self._entry_fn = entry_fn
+        self._context_state = context_state
+        self._context = Context(self._context_state)
+
+        self._state = RunnerState.UNINITIALIZED
+
+        self._input_message_token = None
+
+        # Before we start, we need to convert the input message to the workflow input type
+        self._input_message = input_message
+
+        self._exporter_manager = exporter_manager
+
+    @property
+    def context(self) -> Context:
+        return self._context
+
+    def convert(self, value: typing.Any, to_type: type[_T]) -> _T:
+        return self._entry_fn.convert(value, to_type)
+
+    async def __aenter__(self):
+
+        # Set the input message on the context
+        self._input_message_token = self._context_state.input_message.set(self._input_message)
+
+        # Create reactive event stream
+        self._context_state.event_stream.set(Subject())
+        self._context_state.active_function.set(InvocationNode(
+            function_name="root",
+            function_id="root",
+        ))
+
+        if (self._state == RunnerState.UNINITIALIZED):
+            self._state = RunnerState.INITIALIZED
+        else:
+            raise ValueError("Cannot enter the context more than once")
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+
+        if (self._input_message_token is None):
+            raise ValueError("Cannot exit the context without entering it")
+
+        self._context_state.input_message.reset(self._input_message_token)
+
+        if (self._state not in (RunnerState.COMPLETED, RunnerState.FAILED)):
+            raise ValueError("Cannot exit the context without completing the workflow")
+
+    @typing.overload
+    async def result(self) -> typing.Any:
+        ...
+
+    @typing.overload
+    async def result(self, to_type: type[_T]) -> _T:
+        ...
+
+    async def result(self, to_type: type | None = None):
+
+        if (self._state != RunnerState.INITIALIZED):
+            raise ValueError("Cannot run the workflow without entering the context")
+
+        token_run_id = None
+        token_trace_id = None
+        try:
+            self._state = RunnerState.RUNNING
+
+            if (not self._entry_fn.has_single_output):
+                raise ValueError("Workflow does not support single output")
+
+            # Establish workflow run and trace identifiers
+            existing_run_id = self._context_state.workflow_run_id.get()
+            existing_trace_id = self._context_state.workflow_trace_id.get()
+
+            workflow_run_id = existing_run_id or str(uuid.uuid4())
+
+            workflow_trace_id = existing_trace_id or uuid.uuid4().int
+
+            token_run_id = self._context_state.workflow_run_id.set(workflow_run_id)
+            token_trace_id = self._context_state.workflow_trace_id.set(workflow_trace_id)
+
+            # Prepare workflow-level intermediate step identifiers
+            workflow_step_uuid = str(uuid.uuid4())
+            workflow_name = getattr(self._entry_fn, 'instance_name', None) or "workflow"
+
+            async with self._exporter_manager.start(context_state=self._context_state):
+                # Emit WORKFLOW_START
+                start_metadata = TraceMetadata(
+                    provided_metadata={
+                        "workflow_run_id": workflow_run_id,
+                        "workflow_trace_id": f"{workflow_trace_id:032x}",
+                        "conversation_id": self._context_state.conversation_id.get(),
+                    })
+                self._context.intermediate_step_manager.push_intermediate_step(
+                    IntermediateStepPayload(UUID=workflow_step_uuid,
+                                            event_type=IntermediateStepType.WORKFLOW_START,
+                                            name=workflow_name,
+                                            metadata=start_metadata))
+
+                result = await self._entry_fn.ainvoke(self._input_message, to_type=to_type)  # type: ignore
+
+                # Emit WORKFLOW_END with output
+                end_metadata = TraceMetadata(
+                    provided_metadata={
+                        "workflow_run_id": workflow_run_id,
+                        "workflow_trace_id": f"{workflow_trace_id:032x}",
+                        "conversation_id": self._context_state.conversation_id.get(),
+                    })
+                self._context.intermediate_step_manager.push_intermediate_step(
+                    IntermediateStepPayload(UUID=workflow_step_uuid,
+                                            event_type=IntermediateStepType.WORKFLOW_END,
+                                            name=workflow_name,
+                                            metadata=end_metadata,
+                                            data=StreamEventData(output=result)))
+
+                event_stream = self._context_state.event_stream.get()
+                if event_stream:
+                    event_stream.on_complete()
+
+            self._state = RunnerState.COMPLETED
+
+            return result
+        except Exception as e:
+            logger.error("Error running workflow: %s", e)
+            event_stream = self._context_state.event_stream.get()
+            if event_stream:
+                event_stream.on_complete()
+            self._state = RunnerState.FAILED
+            raise
+        finally:
+            if token_run_id is not None:
+                self._context_state.workflow_run_id.reset(token_run_id)
+            if token_trace_id is not None:
+                self._context_state.workflow_trace_id.reset(token_trace_id)
+
+    async def result_stream(self, to_type: type | None = None):
+
+        if (self._state != RunnerState.INITIALIZED):
+            raise ValueError("Cannot run the workflow without entering the context")
+
+        token_run_id = None
+        token_trace_id = None
+        try:
+            self._state = RunnerState.RUNNING
+
+            if (not self._entry_fn.has_streaming_output):
+                raise ValueError("Workflow does not support streaming output")
+
+            # Establish workflow run and trace identifiers
+            existing_run_id = self._context_state.workflow_run_id.get()
+            existing_trace_id = self._context_state.workflow_trace_id.get()
+
+            workflow_run_id = existing_run_id or str(uuid.uuid4())
+
+            workflow_trace_id = existing_trace_id or uuid.uuid4().int
+
+            token_run_id = self._context_state.workflow_run_id.set(workflow_run_id)
+            token_trace_id = self._context_state.workflow_trace_id.set(workflow_trace_id)
+
+            # Prepare workflow-level intermediate step identifiers
+            workflow_step_uuid = str(uuid.uuid4())
+            workflow_name = getattr(self._entry_fn, 'instance_name', None) or "workflow"
+
+            # Run the workflow
+            async with self._exporter_manager.start(context_state=self._context_state):
+                # Emit WORKFLOW_START
+                start_metadata = TraceMetadata(
+                    provided_metadata={
+                        "workflow_run_id": workflow_run_id,
+                        "workflow_trace_id": f"{workflow_trace_id:032x}",
+                        "conversation_id": self._context_state.conversation_id.get(),
+                    })
+                self._context.intermediate_step_manager.push_intermediate_step(
+                    IntermediateStepPayload(UUID=workflow_step_uuid,
+                                            event_type=IntermediateStepType.WORKFLOW_START,
+                                            name=workflow_name,
+                                            metadata=start_metadata))
+
+                async for m in self._entry_fn.astream(self._input_message, to_type=to_type):  # type: ignore
+                    yield m
+
+                # Emit WORKFLOW_END
+                end_metadata = TraceMetadata(
+                    provided_metadata={
+                        "workflow_run_id": workflow_run_id,
+                        "workflow_trace_id": f"{workflow_trace_id:032x}",
+                        "conversation_id": self._context_state.conversation_id.get(),
+                    })
+                self._context.intermediate_step_manager.push_intermediate_step(
+                    IntermediateStepPayload(UUID=workflow_step_uuid,
+                                            event_type=IntermediateStepType.WORKFLOW_END,
+                                            name=workflow_name,
+                                            metadata=end_metadata))
+                self._state = RunnerState.COMPLETED
+
+                # Close the intermediate stream
+                event_stream = self._context_state.event_stream.get()
+                if event_stream:
+                    event_stream.on_complete()
+
+        except Exception as e:
+            logger.error("Error running workflow: %s", e)
+            event_stream = self._context_state.event_stream.get()
+            if event_stream:
+                event_stream.on_complete()
+            self._state = RunnerState.FAILED
+            raise
+        finally:
+            if token_run_id is not None:
+                self._context_state.workflow_run_id.reset(token_run_id)
+            if token_trace_id is not None:
+                self._context_state.workflow_trace_id.reset(token_trace_id)
+
+
+# Compatibility aliases with previous releases
+AIQRunnerState = RunnerState
+AIQRunner = Runner

nat/runtime/session.py

@@ -0,0 +1,223 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import contextvars
+import typing
+import uuid
+from collections.abc import Awaitable
+from collections.abc import Callable
+from contextlib import asynccontextmanager
+from contextlib import nullcontext
+
+from fastapi import WebSocket
+from starlette.requests import HTTPConnection
+from starlette.requests import Request
+
+from nat.builder.context import Context
+from nat.builder.context import ContextState
+from nat.builder.workflow import Workflow
+from nat.data_models.authentication import AuthenticatedContext
+from nat.data_models.authentication import AuthFlowType
+from nat.data_models.authentication import AuthProviderBaseConfig
+from nat.data_models.config import Config
+from nat.data_models.interactive import HumanResponse
+from nat.data_models.interactive import InteractionPrompt
+
+_T = typing.TypeVar("_T")
+
+
+class UserManagerBase:
+    pass
+
+
+class SessionManager:
+
+    def __init__(self, workflow: Workflow, max_concurrency: int = 8):
+        """
+        The SessionManager class is used to run and manage a user workflow session. It runs and manages the context,
+        and configuration of a workflow with the specified concurrency.
+
+        Parameters
+        ----------
+        workflow : Workflow
+            The workflow to run
+        max_concurrency : int, optional
+            The maximum number of simultaneous workflow invocations, by default 8
+        """
+
+        if (workflow is None):
+            raise ValueError("Workflow cannot be None")
+
+        self._workflow: Workflow = workflow
+
+        self._max_concurrency = max_concurrency
+        self._context_state = ContextState.get()
+        self._context = Context(self._context_state)
+
+        # We save the context because Uvicorn spawns a new process
+        # for each request, and we need to restore the context vars
+        self._saved_context = contextvars.copy_context()
+
+        if (max_concurrency > 0):
+            self._semaphore = asyncio.Semaphore(max_concurrency)
+        else:
+            # If max_concurrency is 0, then we don't need to limit the concurrency but we still need a context
+            self._semaphore = nullcontext()
+
+    @property
+    def config(self) -> Config:
+        return self._workflow.config
+
+    @property
+    def workflow(self) -> Workflow:
+        return self._workflow
+
+    @property
+    def context(self) -> Context:
+        return self._context
+
+    @asynccontextmanager
+    async def session(self,
+                      user_manager=None,
+                      http_connection: HTTPConnection | None = None,
+                      user_message_id: str | None = None,
+                      conversation_id: str | None = None,
+                      user_input_callback: Callable[[InteractionPrompt], Awaitable[HumanResponse]] = None,
+                      user_authentication_callback: Callable[[AuthProviderBaseConfig, AuthFlowType],
+                                                             Awaitable[AuthenticatedContext | None]] = None):
+
+        token_user_input = None
+        if user_input_callback is not None:
+            token_user_input = self._context_state.user_input_callback.set(user_input_callback)
+
+        token_user_manager = None
+        if user_manager is not None:
+            token_user_manager = self._context_state.user_manager.set(user_manager)
+
+        token_user_authentication = None
+        if user_authentication_callback is not None:
+            token_user_authentication = self._context_state.user_auth_callback.set(user_authentication_callback)
+
+        if isinstance(http_connection, WebSocket):
+            self.set_metadata_from_websocket(http_connection, user_message_id, conversation_id)
+
+        if isinstance(http_connection, Request):
+            self.set_metadata_from_http_request(http_connection)
+
+        try:
+            yield self
+        finally:
+            if token_user_manager is not None:
+                self._context_state.user_manager.reset(token_user_manager)
+            if token_user_input is not None:
+                self._context_state.user_input_callback.reset(token_user_input)
+            if token_user_authentication is not None:
+                self._context_state.user_auth_callback.reset(token_user_authentication)
+
+    @asynccontextmanager
+    async def run(self, message):
+        """
+        Start a workflow run
+        """
+        async with self._semaphore:
+            # Apply the saved context
+            for k, v in self._saved_context.items():
+                k.set(v)
+
+            async with self._workflow.run(message) as runner:
+                yield runner
+
+    def set_metadata_from_http_request(self, request: Request) -> None:
+        """
+        Extracts and sets user metadata request attributes from a HTTP request.
+        If request is None, no attributes are set.
+        """
+        self._context.metadata._request.method = getattr(request, "method", None)
+        self._context.metadata._request.url_path = request.url.path
+        self._context.metadata._request.url_port = request.url.port
+        self._context.metadata._request.url_scheme = request.url.scheme
+        self._context.metadata._request.headers = request.headers
+        self._context.metadata._request.query_params = request.query_params
+        self._context.metadata._request.path_params = request.path_params
+        self._context.metadata._request.client_host = request.client.host
+        self._context.metadata._request.client_port = request.client.port
+        self._context.metadata._request.cookies = request.cookies
+
+        if request.headers.get("conversation-id"):
+            self._context_state.conversation_id.set(request.headers["conversation-id"])
+
+        if request.headers.get("user-message-id"):
+            self._context_state.user_message_id.set(request.headers["user-message-id"])
+
+        # W3C Trace Context header: traceparent: 00-<trace-id>-<span-id>-<flags>
+        traceparent = request.headers.get("traceparent")
+        if traceparent:
+            try:
+                parts = traceparent.split("-")
+                if len(parts) >= 4:
+                    trace_id_hex = parts[1]
+                    if len(trace_id_hex) == 32:
+                        trace_id_int = uuid.UUID(trace_id_hex).int
+                        self._context_state.workflow_trace_id.set(trace_id_int)
+            except Exception:
+                pass
+
+        if not self._context_state.workflow_trace_id.get():
+            workflow_trace_id = request.headers.get("workflow-trace-id")
+            if workflow_trace_id:
+                try:
+                    self._context_state.workflow_trace_id.set(uuid.UUID(workflow_trace_id).int)
+                except Exception:
+                    pass
+
+        workflow_run_id = request.headers.get("workflow-run-id")
+        if workflow_run_id:
+            self._context_state.workflow_run_id.set(workflow_run_id)
+
+    def set_metadata_from_websocket(self,
+                                    websocket: WebSocket,
+                                    user_message_id: str | None,
+                                    conversation_id: str | None) -> None:
+        """
+        Extracts and sets user metadata for WebSocket connections.
+        """
+
+        # Extract cookies from WebSocket headers (similar to HTTP request)
+        if websocket and hasattr(websocket, 'scope') and 'headers' in websocket.scope:
+            cookies = {}
+            for header_name, header_value in websocket.scope.get('headers', []):
+                if header_name == b'cookie':
+                    cookie_header = header_value.decode('utf-8')
+                    # Parse cookie header: "name1=value1; name2=value2"
+                    for cookie in cookie_header.split(';'):
+                        cookie = cookie.strip()
+                        if '=' in cookie:
+                            name, value = cookie.split('=', 1)
+                            cookies[name.strip()] = value.strip()
+
+            # Set cookies in metadata (same as HTTP request)
+            self._context.metadata._request.cookies = cookies
+            self._context_state.metadata.set(self._context.metadata)
+
+        if conversation_id is not None:
+            self._context_state.conversation_id.set(conversation_id)
+
+        if user_message_id is not None:
+            self._context_state.user_message_id.set(user_message_id)
+
+
+# Compatibility aliases with previous releases
+AIQSessionManager = SessionManager

nat/runtime/user_metadata.py

@@ -0,0 +1,130 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 starlette.datastructures import Headers
+from starlette.datastructures import QueryParams
+
+from nat.data_models.api_server import Request
+
+
+class RequestAttributes:
+    """
+    The RequestAttributes class is responsible for managing user http and webscoket session
+    metadata. It provides a way to store and expose session attributes to workflow tools.
+    """
+
+    def __init__(self) -> None:
+        self._request: Request = Request()
+
+    @property
+    def method(self) -> str | None:
+        """
+        This property retrieves the HTTP method of the request.
+        It can be GET, POST, PUT, DELETE, etc.
+
+        Returns:
+            str | None
+        """
+        return self._request.method
+
+    @property
+    def url_path(self) -> str | None:
+        """
+        This property retrieves the path from the URL of the request.
+
+        Returns:
+            str | None
+        """
+        return self._request.url_path
+
+    @property
+    def url_port(self) -> int | None:
+        """
+        This property retrieves the port number from the URL of the request.
+
+        Returns:
+            int | None
+        """
+        return self._request.url_port
+
+    @property
+    def url_scheme(self) -> str | None:
+        """
+        This property retrieves the scheme from the URL of the request.
+
+        Returns:
+            str | None
+        """
+        return self._request.url_scheme
+
+    @property
+    def headers(self) -> Headers | None:
+        """
+        This property retrieves the headers from the request stored in a dictionary-like object.
+
+        Returns:
+            Headers | None
+        """
+        return self._request.headers
+
+    @property
+    def query_params(self) -> QueryParams | None:
+        """
+        This property retrieves the query parameters from the request stored in a dictionary-like object.
+
+        Returns:
+            QueryParams | None
+        """
+        return self._request.query_params
+
+    @property
+    def path_params(self) -> dict[str, str] | None:
+        """
+        This property retrieves the path parameters from the request stored in a dictionary-like object.
+
+        Returns:
+            dict[str, str] | None
+        """
+        return self._request.path_params
+
+    @property
+    def client_host(self) -> str | None:
+        """
+        This property retrieves the clients remote hostname or IP address.
+
+        Returns:
+            str | None
+        """
+        return self._request.client_host
+
+    @property
+    def client_port(self) -> int | None:
+        """
+        This property retrieves the clients remote port number from which the client is connecting to.
+
+        Returns:
+            int | None
+        """
+        return self._request.client_port
+
+    @property
+    def cookies(self) -> dict[str, str] | None:
+        """
+        This property retrieves the cookies from the request stored in a dictionary-like object.
+
+        Returns:
+            dict[str, str] | None
+        """
+        return self._request.cookies