nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.4.0a20251112__py3-none-any.whl

nat/runtime/runner.py

@@ -15,11 +15,16 @@
 
 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
@@ -130,17 +135,60 @@ class Runner:
         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):
-                # Run the workflow
-                result = await self._entry_fn.ainvoke(self._input_message, to_type=to_type)
+                # 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,
+                                            data=StreamEventData(input=self._input_message)))
+
+                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)))
 
-                # Close the intermediate stream
                 event_stream = self._context_state.event_stream.get()
                 if event_stream:
                     event_stream.on_complete()
@@ -154,25 +202,78 @@ class Runner:
             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):
-                async for m in self._entry_fn.astream(self._input_message, to_type=to_type):
+                # 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,
+                                            data=StreamEventData(input=self._input_message)))
+
+                # Collect preview of streaming results for the WORKFLOW_END event
+                output_preview = []
+
+                async for m in self._entry_fn.astream(self._input_message, to_type=to_type):  # type: ignore
+                    if len(output_preview) < 50:
+                        output_preview.append(m)
                     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,
+                                            data=StreamEventData(output=output_preview)))
                 self._state = RunnerState.COMPLETED
 
                 # Close the intermediate stream
@@ -186,8 +287,12 @@ class Runner:
             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

nat/runtime/session.py

@@ -16,6 +16,7 @@
 import asyncio
 import contextvars
 import typing
+import uuid
 from collections.abc import Awaitable
 from collections.abc import Callable
 from contextlib import asynccontextmanager
@@ -111,7 +112,7 @@ class SessionManager:
             token_user_authentication = self._context_state.user_auth_callback.set(user_authentication_callback)
 
         if isinstance(http_connection, WebSocket):
-            self.set_metadata_from_websocket(user_message_id, conversation_id)
+            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)
@@ -161,11 +162,56 @@ class SessionManager:
         if request.headers.get("user-message-id"):
             self._context_state.user_message_id.set(request.headers["user-message-id"])
 
-    def set_metadata_from_websocket(self, user_message_id: str | None, conversation_id: str | None) -> None:
+        # 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.
+        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)
 

nat/settings/global_settings.py

@@ -124,7 +124,7 @@ class Settings(HashableBaseModel):
                 if (short_names[key.local_name] == 1):
                     type_list.append((key.local_name, key.config_type))
 
-            return typing.Union[tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
+            return typing.Union[*tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
 
         RegistryHandlerAnnotation = dict[
             str,
@@ -169,7 +169,7 @@ class Settings(HashableBaseModel):
         if (not os.path.exists(configuration_file)):
             loaded_config = {}
         else:
-            with open(file_path, mode="r", encoding="utf-8") as f:
+            with open(file_path, encoding="utf-8") as f:
                 try:
                     loaded_config = json.load(f)
                 except Exception as e:

nat/tool/chat_completion.py

@@ -63,7 +63,10 @@ async def register_chat_completion(config: ChatCompletionConfig, builder: Builde
             # Generate response using the LLM
             response = await llm.ainvoke(prompt)
 
-            return response
+            if isinstance(response, str):
+                return response
+
+            return response.text()
 
         except Exception as e:
             # Fallback response if LLM call fails

nat/tool/code_execution/code_sandbox.py

@@ -92,7 +92,9 @@ class Sandbox(abc.ABC):
             raise ValueError(f"Language {language} not supported")
 
         generated_code = generated_code.strip().strip("`")
-        code_to_execute = textwrap.dedent("""
+        # Use json.dumps to properly escape the generated_code instead of repr()
+        escaped_code = json.dumps(generated_code)
+        code_to_execute = textwrap.dedent(f"""
             import traceback
             import json
             import os
@@ -101,11 +103,6 @@ class Sandbox(abc.ABC):
             import io
             warnings.filterwarnings('ignore')
             os.environ['OPENBLAS_NUM_THREADS'] = '16'
-        """).strip()
-
-        # Use json.dumps to properly escape the generated_code instead of repr()
-        escaped_code = json.dumps(generated_code)
-        code_to_execute += textwrap.dedent(f"""
 
             generated_code = {escaped_code}
 

nat/tool/code_execution/local_sandbox/Dockerfile.sandbox

@@ -12,43 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Use the base image with Python 3.10 and Flask
-FROM tiangolo/uwsgi-nginx-flask:python3.10
-
-# Install dependencies required for Lean 4 and other tools
-RUN apt-get update && \
-    apt-get install -y curl git && \
-    curl https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -sSf | sh -s -- -y && \
-    /root/.elan/bin/elan toolchain install leanprover/lean4:v4.12.0 && \
-    /root/.elan/bin/elan default leanprover/lean4:v4.12.0 && \
-    /root/.elan/bin/elan self update
-
-# Set environment variables to include Lean and elan/lake in the PATH
-ENV PATH="/root/.elan/bin:$PATH"
-
-# Create Lean project directory and initialize a new Lean project with Mathlib4
-RUN mkdir -p /lean4 && cd /lean4 && \
-    /root/.elan/bin/lake new my_project && \
-    cd my_project && \
-    echo 'leanprover/lean4:v4.12.0' > lean-toolchain && \
-    echo 'require mathlib from git "https://github.com/leanprover-community/mathlib4" @ "v4.12.0"' >> lakefile.lean
-
-# Download and cache Mathlib4 to avoid recompiling, then build the project
-RUN cd /lean4/my_project && \
-    /root/.elan/bin/lake exe cache get && \
-    /root/.elan/bin/lake build
-
-# Set environment variables to include Lean project path
-ENV LEAN_PATH="/lean4/my_project"
-ENV PATH="/lean4/my_project:$PATH"
+# UWSGI_CHEAPER sets the number of initial uWSGI worker processes
+# UWSGI_PROCESSES sets the maximum number of uWSGI worker processes
+ARG UWSGI_CHEAPER=5
+ARG UWSGI_PROCESSES=10
+
+# Use the base image with Python 3.13
+FROM python:3.13-slim-bookworm
+
+
+RUN apt update && \
+    apt upgrade && \
+    apt install -y --no-install-recommends libexpat1 && \
+    apt clean && \
+    rm -rf /var/lib/apt/lists/*
 
 # Set up application code and install Python dependencies
 COPY sandbox.requirements.txt /app/requirements.txt
 RUN pip install --no-cache-dir -r /app/requirements.txt
 COPY local_sandbox_server.py /app/main.py
-
-# Set the working directory to /app
-WORKDIR /app
+RUN mkdir /workspace
 
 # Set Flask app environment variables and ports
 ARG UWSGI_CHEAPER
@@ -58,3 +41,7 @@ ARG UWSGI_PROCESSES
 ENV UWSGI_PROCESSES=$UWSGI_PROCESSES
 
 ENV LISTEN_PORT=6000
+EXPOSE 6000
+
+WORKDIR /app
+CMD uwsgi --http 0.0.0.0:${LISTEN_PORT} --master -p ${UWSGI_PROCESSES} --force-cwd /workspace -w main:app

nat/tool/code_execution/local_sandbox/local_sandbox_server.py

@@ -62,7 +62,7 @@ class CodeExecutionResponse(Response):
         super().__init__(status=status_code, mimetype="application/json", response=result.model_dump_json())
 
     @classmethod
-    def with_error(cls, status_code: int, error_message: str) -> 'CodeExecutionResponse':
+    def with_error(cls, status_code: int, error_message: str) -> CodeExecutionResponse:
         return cls(status_code,
                    CodeExecutionResult(process_status=CodeExecutionStatus.ERROR, stdout="", stderr=error_message))
 
@@ -194,5 +194,10 @@ def execute():
     return do_execute(request)
 
 
+@app.route("/", methods=["GET"])
+def status() -> tuple[dict[str, str], int]:
+    return ({"status": "ok"}, 200)
+
+
 if __name__ == '__main__':
     app.run(port=6000)

nat/tool/code_execution/local_sandbox/sandbox.requirements.txt

@@ -1,6 +1,8 @@
+Flask==3.1
 numpy
 pandas
 scipy
 ipython
 plotly
 pydantic
+pyuwsgi==2.0.*

nat/tool/code_execution/local_sandbox/start_local_sandbox.sh

@@ -19,7 +19,11 @@
 
 DOCKER_COMMAND=${DOCKER_COMMAND:-"docker"}
 SANDBOX_NAME=${1:-'local-sandbox'}
-NUM_THREADS=10
+
+# UWSGI_CHEAPER sets the number of initial uWSGI worker processes
+# UWSGI_PROCESSES sets the maximum number of uWSGI worker processes
+UWSGI_CHEAPER=${UWSGI_CHEAPER:-5}
+UWSGI_PROCESSES=${UWSGI_PROCESSES:-10}
 
 # Get the output_data directory path for mounting
 # Priority: command line argument > environment variable > default path (current directory)
@@ -37,14 +41,16 @@ fi
 # Check if the Docker image already exists
 if ! ${DOCKER_COMMAND} images ${SANDBOX_NAME} | grep -q "${SANDBOX_NAME}"; then
     echo "Docker image not found locally. Building ${SANDBOX_NAME}..."
-    ${DOCKER_COMMAND} build --tag=${SANDBOX_NAME} --build-arg="UWSGI_PROCESSES=$((${NUM_THREADS} * 10))" --build-arg="UWSGI_CHEAPER=${NUM_THREADS}" -f Dockerfile.sandbox .
+    ${DOCKER_COMMAND} build --tag=${SANDBOX_NAME} \
+        --build-arg="UWSGI_PROCESSES=${UWSGI_PROCESSES}" \
+        --build-arg="UWSGI_CHEAPER=${UWSGI_CHEAPER}" \
+        -f Dockerfile.sandbox .
 else
     echo "Using existing Docker image: ${SANDBOX_NAME}"
 fi
 
 # Mount the output_data directory directly so files created in container appear in the local directory
-${DOCKER_COMMAND} run --rm --name=local-sandbox \
+${DOCKER_COMMAND} run --rm -ti --name=local-sandbox \
   --network=host \
   -v "${OUTPUT_DATA_PATH}:/workspace" \
-  -w /workspace \
   ${SANDBOX_NAME}

nat/tool/datetime_tools.py

@@ -72,7 +72,7 @@ async def current_datetime(_config: CurrentTimeToolConfig, _builder: Builder):
         timezone_obj = _get_timezone_obj(headers)
 
         now = datetime.datetime.now(timezone_obj)
-        now_machine_readable = now.strftime(("%Y-%m-%d %H:%M:%S %z"))
+        now_machine_readable = now.strftime("%Y-%m-%d %H:%M:%S %z")
 
         # Returns the current time in machine readable format with timezone offset.
         return f"The current time of day is {now_machine_readable}"