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

nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py

@@ -689,10 +689,13 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def post_openai_api_compatible(response: Response, request: Request, payload: request_type):
                 # Check if streaming is requested
+
+                response.headers["Content-Type"] = "application/json"
                 stream_requested = getattr(payload, 'stream', False)
 
                 async with session_manager.session(http_connection=request):
                     if stream_requested:
+
                         # Return streaming response
                         return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
                                                  content=generate_streaming_response_as_str(
@@ -703,40 +706,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                                                      result_type=ChatResponseChunk,
                                                      output_type=ChatResponseChunk))
 
-                    # Return single response - check if workflow supports non-streaming
-                    try:
-                        response.headers["Content-Type"] = "application/json"
-                        return await generate_single_response(payload, session_manager, result_type=ChatResponse)
-                    except ValueError as e:
-                        if "Cannot get a single output value for streaming workflows" in str(e):
-                            # Workflow only supports streaming, but client requested non-streaming
-                            # Fall back to streaming and collect the result
-                            chunks = []
-                            async for chunk_str in generate_streaming_response_as_str(
-                                    payload,
-                                    session_manager=session_manager,
-                                    streaming=True,
-                                    step_adaptor=self.get_step_adaptor(),
-                                    result_type=ChatResponseChunk,
-                                    output_type=ChatResponseChunk):
-                                if chunk_str.startswith("data: ") and not chunk_str.startswith("data: [DONE]"):
-                                    chunk_data = chunk_str[6:].strip()  # Remove "data: " prefix
-                                    if chunk_data:
-                                        try:
-                                            chunk_json = ChatResponseChunk.model_validate_json(chunk_data)
-                                            if (chunk_json.choices and len(chunk_json.choices) > 0
-                                                    and chunk_json.choices[0].delta
-                                                    and chunk_json.choices[0].delta.content is not None):
-                                                chunks.append(chunk_json.choices[0].delta.content)
-                                        except Exception:
-                                            continue
-
-                            # Create a single response from collected chunks
-                            content = "".join(chunks)
-                            single_response = ChatResponse.from_string(content)
-                            response.headers["Content-Type"] = "application/json"
-                            return single_response
-                        raise
+                    return await generate_single_response(payload, session_manager, result_type=ChatResponse)
 
             return post_openai_api_compatible
 
@@ -1128,7 +1098,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                     if configured_group.config.type != "mcp_client":
                         continue
 
-                    from nat.plugins.mcp.client_impl import MCPClientConfig
+                    from nat.plugins.mcp.client_config import MCPClientConfig
 
                     config = configured_group.config
                     assert isinstance(config, MCPClientConfig)

nat/front_ends/fastapi/message_validator.py

@@ -139,8 +139,10 @@ class MessageValidator:
                     text_content: str = str(data_model.payload)
                 validated_message_content = SystemResponseContent(text=text_content)
 
-            elif (isinstance(data_model, ChatResponse | ChatResponseChunk)):
+            elif isinstance(data_model, ChatResponse):
                 validated_message_content = SystemResponseContent(text=data_model.choices[0].message.content)
+            elif isinstance(data_model, ChatResponseChunk):
+                validated_message_content = SystemResponseContent(text=data_model.choices[0].delta.content)
 
             elif (isinstance(data_model, ResponseIntermediateStep)):
                 validated_message_content = SystemIntermediateStepContent(name=data_model.name,

nat/observability/exporter/span_exporter.py

@@ -126,6 +126,7 @@ class SpanExporter(ProcessingExporter[InputSpanT, OutputSpanT], SerializeMixin):
 
         parent_span = None
         span_ctx = None
+        workflow_trace_id = self._context_state.workflow_trace_id.get()
 
         # Look up the parent span to establish hierarchy
         # event.parent_id is the UUID of the last START step with a different UUID from current step
@@ -141,6 +142,9 @@ class SpanExporter(ProcessingExporter[InputSpanT, OutputSpanT], SerializeMixin):
             parent_span = parent_span.model_copy() if isinstance(parent_span, Span) else None
             if parent_span and parent_span.context:
                 span_ctx = SpanContext(trace_id=parent_span.context.trace_id)
+        # No parent: adopt workflow trace id if available to keep all spans in the same trace
+        if span_ctx is None and workflow_trace_id:
+            span_ctx = SpanContext(trace_id=workflow_trace_id)
 
         # Extract start/end times from the step
         # By convention, `span_event_timestamp` is the time we started, `event_timestamp` is the time we ended.
@@ -154,23 +158,39 @@ class SpanExporter(ProcessingExporter[InputSpanT, OutputSpanT], SerializeMixin):
         else:
             sub_span_name = f"{event.payload.event_type}"
 
+        # Prefer parent/context trace id for attribute, else workflow trace id
+        _attr_trace_id = None
+        if span_ctx is not None:
+            _attr_trace_id = span_ctx.trace_id
+        elif parent_span and parent_span.context:
+            _attr_trace_id = parent_span.context.trace_id
+        elif workflow_trace_id:
+            _attr_trace_id = workflow_trace_id
+
+        attributes = {
+            f"{self._span_prefix}.event_type":
+                event.payload.event_type.value,
+            f"{self._span_prefix}.function.id":
+                event.function_ancestry.function_id if event.function_ancestry else "unknown",
+            f"{self._span_prefix}.function.name":
+                event.function_ancestry.function_name if event.function_ancestry else "unknown",
+            f"{self._span_prefix}.subspan.name":
+                event.payload.name or "",
+            f"{self._span_prefix}.event_timestamp":
+                event.event_timestamp,
+            f"{self._span_prefix}.framework":
+                event.payload.framework.value if event.payload.framework else "unknown",
+            f"{self._span_prefix}.conversation.id":
+                self._context_state.conversation_id.get() or "unknown",
+            f"{self._span_prefix}.workflow.run_id":
+                self._context_state.workflow_run_id.get() or "unknown",
+            f"{self._span_prefix}.workflow.trace_id": (f"{_attr_trace_id:032x}" if _attr_trace_id else "unknown"),
+        }
+
         sub_span = Span(name=sub_span_name,
                         parent=parent_span,
                         context=span_ctx,
-                        attributes={
-                            f"{self._span_prefix}.event_type":
-                                event.payload.event_type.value,
-                            f"{self._span_prefix}.function.id":
-                                event.function_ancestry.function_id if event.function_ancestry else "unknown",
-                            f"{self._span_prefix}.function.name":
-                                event.function_ancestry.function_name if event.function_ancestry else "unknown",
-                            f"{self._span_prefix}.subspan.name":
-                                event.payload.name or "",
-                            f"{self._span_prefix}.event_timestamp":
-                                event.event_timestamp,
-                            f"{self._span_prefix}.framework":
-                                event.payload.framework.value if event.payload.framework else "unknown",
-                        },
+                        attributes=attributes,
                         start_time=start_ns)
 
         span_kind = event_type_to_span_kind(event.event_type)

nat/observability/register.py

@@ -77,6 +77,14 @@ async def console_logging_method(config: ConsoleLoggingMethodConfig, builder: Bu
     level = getattr(logging, config.level.upper(), logging.INFO)
     handler = logging.StreamHandler(stream=sys.stdout)
     handler.setLevel(level)
+
+    # Set formatter to match the default CLI format
+    formatter = logging.Formatter(
+        fmt="%(asctime)s - %(levelname)-8s - %(name)s:%(lineno)d - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    handler.setFormatter(formatter)
+
     yield handler
 
 
@@ -95,4 +103,12 @@ async def file_logging_method(config: FileLoggingMethod, builder: Builder):
     level = getattr(logging, config.level.upper(), logging.INFO)
     handler = logging.FileHandler(filename=config.path, mode="a", encoding="utf-8")
     handler.setLevel(level)
+
+    # Set formatter to match the default CLI format
+    formatter = logging.Formatter(
+        fmt="%(asctime)s - %(levelname)-8s - %(name)s:%(lineno)d - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    handler.setFormatter(formatter)
+
     yield handler

nat/profiler/decorators/framework_wrapper.py

@@ -123,7 +123,7 @@ def set_framework_profiler_handler(
                 except ImportError as e:
                     logger.warning(
                         "ADK profiler not available. " +
-                        "Install NAT with ADK extras: pip install 'nvidia-nat[adk]'. Error: %s",
+                        "Install NAT with ADK extras: pip install \"nvidia-nat[adk]\". Error: %s",
                         e)
                 else:
                     handler = ADKProfilerHandler()

nat/profiler/forecasting/models/linear_model.py

@@ -36,7 +36,7 @@ class LinearModel(ForecastingBaseModel):
         except ImportError:
             logger.error(
                 "scikit-learn is not installed. Please install scikit-learn to use the LinearModel "
-                "profiling model or install `nvidia-nat[profiler]` to install all necessary profiling packages.")
+                "profiling model or install \"nvidia-nat[profiler]\" to install all necessary profiling packages.")
 
             raise
 

nat/profiler/forecasting/models/random_forest_regressor.py

@@ -36,7 +36,7 @@ class RandomForestModel(ForecastingBaseModel):
         except ImportError:
             logger.error(
                 "scikit-learn is not installed. Please install scikit-learn to use the RandomForest "
-                "profiling model or install `nvidia-nat[profiler]` to install all necessary profiling packages.")
+                "profiling model or install \"nvidia-nat[profiler]\" to install all necessary profiling packages.")
 
             raise
 

nat/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py

@@ -304,7 +304,7 @@ def save_gantt_chart(all_nodes: list[CallNode], output_path: str) -> None:
         import matplotlib.pyplot as plt
     except ImportError:
         logger.error("matplotlib is not installed. Please install matplotlib to use generate plots for the profiler "
-                     "or install `nvidia-nat[profiler]` to install all necessary profiling packages.")
+                     "or install \"nvidia-nat[profiler]\" to install all necessary profiling packages.")
 
         raise
 

nat/profiler/inference_optimization/experimental/prefix_span_analysis.py

@@ -212,7 +212,7 @@ def run_prefixspan(sequences_map: dict[int, list[PrefixCallNode]],
         from prefixspan import PrefixSpan
     except ImportError:
         logger.error("prefixspan is not installed. Please install prefixspan to run the prefix analysis in the "
-                     "profiler or install `nvidia-nat[profiler]` to install all necessary profiling packages.")
+                     "profiler or install \"nvidia-nat[profiler]\" to install all necessary profiling packages.")
 
         raise
 

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,59 @@ 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))
+
+                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()
@@ -155,25 +202,71 @@ 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))
+
+                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
@@ -187,8 +280,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
@@ -161,12 +162,37 @@ class SessionManager:
         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.
+        Extracts and sets user metadata for WebSocket connections.
         """
 
         # Extract cookies from WebSocket headers (similar to HTTP request)

nat/tool/memory_tools/add_memory_tool.py

@@ -30,10 +30,10 @@ logger = logging.getLogger(__name__)
 class AddToolConfig(FunctionBaseConfig, name="add_memory"):
     """Function to add memory to a hosted memory platform."""
 
-    description: str = Field(default=("Tool to add memory about a user's interactions to a system "
+    description: str = Field(default=("Tool to add a memory about a user's interactions to a system "
                                       "for retrieval later."),
                              description="The description of this function's use for tool calling agents.")
-    memory: MemoryRef = Field(default="saas_memory",
+    memory: MemoryRef = Field(default=MemoryRef("saas_memory"),
                               description=("Instance name of the memory client instance from the workflow "
                                            "configuration object."))
 
@@ -46,7 +46,7 @@ async def add_memory_tool(config: AddToolConfig, builder: Builder):
     from langchain_core.tools import ToolException
 
     # First, retrieve the memory client
-    memory_editor = builder.get_memory_client(config.memory)
+    memory_editor = await builder.get_memory_client(config.memory)
 
     async def _arun(item: MemoryItem) -> str:
         """

nat/tool/memory_tools/delete_memory_tool.py

@@ -30,10 +30,9 @@ logger = logging.getLogger(__name__)
 class DeleteToolConfig(FunctionBaseConfig, name="delete_memory"):
     """Function to delete memory from a hosted memory platform."""
 
-    description: str = Field(default=("Tool to retrieve memory about a user's "
-                                      "interactions to help answer questions in a personalized way."),
+    description: str = Field(default="Tool to delete a memory from a hosted memory platform.",
                              description="The description of this function's use for tool calling agents.")
-    memory: MemoryRef = Field(default="saas_memory",
+    memory: MemoryRef = Field(default=MemoryRef("saas_memory"),
                               description=("Instance name of the memory client instance from the workflow "
                                            "configuration object."))
 
@@ -47,7 +46,7 @@ async def delete_memory_tool(config: DeleteToolConfig, builder: Builder):
     from langchain_core.tools import ToolException
 
     # First, retrieve the memory client
-    memory_editor = builder.get_memory_client(config.memory)
+    memory_editor = await builder.get_memory_client(config.memory)
 
     async def _arun(user_id: str) -> str:
         """

nat/tool/memory_tools/get_memory_tool.py

@@ -30,10 +30,10 @@ logger = logging.getLogger(__name__)
 class GetToolConfig(FunctionBaseConfig, name="get_memory"):
     """Function to get memory to a hosted memory platform."""
 
-    description: str = Field(default=("Tool to retrieve memory about a user's "
+    description: str = Field(default=("Tool to retrieve a memory about a user's "
                                       "interactions to help answer questions in a personalized way."),
                              description="The description of this function's use for tool calling agents.")
-    memory: MemoryRef = Field(default="saas_memory",
+    memory: MemoryRef = Field(default=MemoryRef("saas_memory"),
                               description=("Instance name of the memory client instance from the workflow "
                                            "configuration object."))
 
@@ -49,7 +49,7 @@ async def get_memory_tool(config: GetToolConfig, builder: Builder):
     from langchain_core.tools import ToolException
 
     # First, retrieve the memory client
-    memory_editor = builder.get_memory_client(config.memory)
+    memory_editor = await builder.get_memory_client(config.memory)
 
     async def _arun(search_input: SearchMemoryInput) -> str:
         """
@@ -67,6 +67,6 @@ async def get_memory_tool(config: GetToolConfig, builder: Builder):
 
         except Exception as e:
 
-            raise ToolException(f"Error retreiving memory: {e}") from e
+            raise ToolException(f"Error retrieving memory: {e}") from e
 
     yield FunctionInfo.from_fn(_arun, description=config.description)