nvidia-nat 1.4.0a20251112__py3-none-any.whl → 1.4.0a20260113__py3-none-any.whl

nat/front_ends/fastapi/job_store.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -208,13 +208,14 @@ class JobStore(DaskClientMixin):
         AsyncSession
             An active SQLAlchemy async session with an open transaction.
         """
-        async with self._session() as session:
-            async with session.begin():
-                yield session
-
-        # Removes the current task key from the session registry, preventing
-        # potential memory leaks
-        await self._session.remove()
+        try:
+            async with self._session() as session:
+                async with session.begin():
+                    yield session
+        finally:
+            # Removes the current task key from the session registry, preventing
+            # potential memory leaks
+            await self._session.remove()
 
     def ensure_job_id(self, job_id: str | None) -> str:
         """
@@ -502,13 +503,14 @@ class JobStore(DaskClientMixin):
 
         return updated_at + timedelta(seconds=job.expiry_seconds)
 
-    async def cleanup_expired_jobs(self):
+    async def cleanup_expired_jobs(self) -> int:
         """
         Cleanup expired jobs, keeping the most recent one.
 
         Updated_at is used instead of created_at to determine the most recent job. This is because jobs may not be
         processed in the order they are created.
         """
+        logger.info("Starting cleanup of expired jobs")
         now = datetime.now(UTC)
 
         stmt = select(JobInfo).where(
@@ -536,9 +538,11 @@ class JobStore(DaskClientMixin):
                         elif os.path.isdir(job.output_path):
                             shutil.rmtree(job.output_path)
 
-            if len(expired_ids) > 0:
+            num_expired = len(expired_ids)
+            if num_expired > 0:
                 successfully_expired = []
                 for job_id in expired_ids:
+                    var = None
                     try:
                         var = Variable(name=job_id, client=client)
                         try:
@@ -549,14 +553,22 @@ class JobStore(DaskClientMixin):
                         except TimeoutError:
                             pass
 
-                        var.delete()
                         successfully_expired.append(job_id)
                     except Exception:
                         logger.exception("Failed to expire %s", job_id)
 
+                    finally:
+                        if var is not None:
+                            try:
+                                var.delete()
+                            except Exception:
+                                logger.exception("Failed to delete variable %s", job_id)
+
                 await session.execute(
                     update(JobInfo).where(JobInfo.job_id.in_(successfully_expired)).values(is_expired=True))
 
+            return num_expired
+
 
 def get_db_engine(db_url: str | None = None, echo: bool = False, use_async: bool = True) -> "Engine | AsyncEngine":
     """

nat/front_ends/fastapi/main.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/fastapi/message_handler.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -30,6 +30,7 @@ from nat.data_models.api_server import ChatResponse
 from nat.data_models.api_server import ChatResponseChunk
 from nat.data_models.api_server import Error
 from nat.data_models.api_server import ErrorTypes
+from nat.data_models.api_server import ResponseObservabilityTrace
 from nat.data_models.api_server import ResponsePayloadOutput
 from nat.data_models.api_server import ResponseSerializable
 from nat.data_models.api_server import SystemResponseContent
@@ -38,6 +39,7 @@ from nat.data_models.api_server import UserMessageContentRoleType
 from nat.data_models.api_server import UserMessages
 from nat.data_models.api_server import WebSocketMessageStatus
 from nat.data_models.api_server import WebSocketMessageType
+from nat.data_models.api_server import WebSocketObservabilityTraceMessage
 from nat.data_models.api_server import WebSocketSystemInteractionMessage
 from nat.data_models.api_server import WebSocketSystemIntermediateStepMessage
 from nat.data_models.api_server import WebSocketSystemResponseTokenMessage
@@ -69,14 +71,15 @@ class WebSocketMessageHandler:
         self._conversation_id: str | None = None
         self._workflow_schema_type: str | None = None
         self._user_interaction_response: asyncio.Future[TextContent] | None = None
+        self._pending_observability_trace: ResponseObservabilityTrace | None = None
 
         self._flow_handler: FlowHandlerBase | None = None
 
         self._schema_output_mapping: dict[str, type[BaseModel] | type[None]] = {
-            WorkflowSchemaType.GENERATE: self._session_manager.workflow.single_output_schema,
+            WorkflowSchemaType.GENERATE: self._session_manager.get_workflow_single_output_schema(),
             WorkflowSchemaType.CHAT: ChatResponse,
             WorkflowSchemaType.CHAT_STREAM: ChatResponseChunk,
-            WorkflowSchemaType.GENERATE_STREAM: self._session_manager.workflow.streaming_output_schema,
+            WorkflowSchemaType.GENERATE_STREAM: self._session_manager.get_workflow_streaming_output_schema(),
         }
 
     def set_flow_handler(self, flow_handler: FlowHandlerBase) -> None:
@@ -175,6 +178,7 @@ class WebSocketMessageHandler:
             self._message_parent_id = user_message_as_validated_type.id
             self._workflow_schema_type = user_message_as_validated_type.schema_type
             self._conversation_id = user_message_as_validated_type.conversation_id
+            self._pending_observability_trace = None
 
             message_content: typing.Any = await self._process_websocket_user_message(user_message_as_validated_type)
 
@@ -250,6 +254,13 @@ class WebSocketMessageHandler:
                     content=content,
                     status=status)
 
+            elif issubclass(message_schema, WebSocketObservabilityTraceMessage):
+                message = await self._message_validator.create_observability_trace_message(
+                    message_id=message_id,
+                    parent_id=self._message_parent_id,
+                    conversation_id=self._conversation_id,
+                    content=content)
+
             elif isinstance(content, Error):
                 raise ValidationError(f"Invalid input data creating websocket message. {data_model.model_dump_json()}")
 
@@ -327,12 +338,18 @@ class WebSocketMessageHandler:
                                                      user_authentication_callback=auth_callback) as session:
 
                 async for value in generate_streaming_response(payload,
-                                                               session_manager=session,
+                                                               session=session,
                                                                streaming=True,
                                                                step_adaptor=self._step_adaptor,
                                                                result_type=result_type,
                                                                output_type=output_type):
 
+                    # Store observability trace to send after completion message
+                    if isinstance(value, ResponseObservabilityTrace):
+                        if self._pending_observability_trace is None:
+                            self._pending_observability_trace = value
+                        continue
+
                     if not isinstance(value, ResponseSerializable):
                         value = ResponsePayloadOutput(payload=value)
 
@@ -342,3 +359,9 @@ class WebSocketMessageHandler:
             await self.create_websocket_message(data_model=SystemResponseContent(),
                                                 message_type=WebSocketMessageType.RESPONSE_MESSAGE,
                                                 status=WebSocketMessageStatus.COMPLETE)
+
+            # Send observability trace after completion message
+            if self._pending_observability_trace is not None:
+                await self.create_websocket_message(data_model=self._pending_observability_trace,
+                                                    message_type=WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE)
+                self._pending_observability_trace = None

nat/front_ends/fastapi/message_validator.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -27,13 +27,16 @@ from nat.data_models.api_server import ChatResponse
 from nat.data_models.api_server import ChatResponseChunk
 from nat.data_models.api_server import Error
 from nat.data_models.api_server import ErrorTypes
+from nat.data_models.api_server import ObservabilityTraceContent
 from nat.data_models.api_server import ResponseIntermediateStep
+from nat.data_models.api_server import ResponseObservabilityTrace
 from nat.data_models.api_server import ResponsePayloadOutput
 from nat.data_models.api_server import SystemIntermediateStepContent
 from nat.data_models.api_server import SystemResponseContent
 from nat.data_models.api_server import TextContent
 from nat.data_models.api_server import WebSocketMessageStatus
 from nat.data_models.api_server import WebSocketMessageType
+from nat.data_models.api_server import WebSocketObservabilityTraceMessage
 from nat.data_models.api_server import WebSocketSystemInteractionMessage
 from nat.data_models.api_server import WebSocketSystemIntermediateStepMessage
 from nat.data_models.api_server import WebSocketSystemResponseTokenMessage
@@ -67,11 +70,24 @@ class MessageValidator:
             WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE: WebSocketSystemIntermediateStepMessage,
             WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE: WebSocketSystemInteractionMessage,
             WebSocketMessageType.USER_INTERACTION_MESSAGE: WebSocketUserInteractionResponseMessage,
-            WebSocketMessageType.ERROR_MESSAGE: Error
+            WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE: WebSocketObservabilityTraceMessage,
+            WebSocketMessageType.ERROR_MESSAGE: Error,
         }
 
         self._message_parent_id: str = "default_id"
 
+    def _get_observability_trace_id_from_context(self) -> str | None:
+        """
+        Retrieves observability_trace_id from Context
+
+        :return: observability_trace_id if available, None otherwise.
+        """
+        try:
+            from nat.builder.context import Context
+            return Context.get().observability_trace_id
+        except (ImportError, AttributeError, KeyError):
+            return None
+
     async def validate_message(self, message: dict[str, Any]) -> BaseModel:
         """
         Validates an incoming WebSocket message against its expected schema.
@@ -147,6 +163,9 @@ class MessageValidator:
             elif (isinstance(data_model, ResponseIntermediateStep)):
                 validated_message_content = SystemIntermediateStepContent(name=data_model.name,
                                                                           payload=data_model.payload)
+            elif (isinstance(data_model, ResponseObservabilityTrace)):
+                validated_message_content = ObservabilityTraceContent(
+                    observability_trace_id=data_model.observability_trace_id)
             elif (isinstance(data_model, HumanPromptBase)):
                 validated_message_content = data_model
             elif (isinstance(data_model, SystemResponseContent)):
@@ -212,6 +231,9 @@ class MessageValidator:
             elif (isinstance(data_model, ResponseIntermediateStep)):
                 validated_message_type = WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE
 
+            elif (isinstance(data_model, ResponseObservabilityTrace)):
+                validated_message_type = WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE
+
             elif (isinstance(data_model, HumanPromptBase)):
                 validated_message_type = WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE
             else:
@@ -349,3 +371,33 @@ class MessageValidator:
         except Exception as e:
             logger.exception("Error creating system interaction message: %s", str(e))
             return None
+
+    async def create_observability_trace_message(
+        self,
+        *,
+        message_id: str | None = str(uuid.uuid4()),
+        parent_id: str = "default",
+        conversation_id: str | None = None,
+        content: ObservabilityTraceContent,
+        timestamp: str = str(datetime.datetime.now(datetime.UTC))
+    ) -> WebSocketObservabilityTraceMessage | None:
+        """
+        Creates an observability trace message.
+
+        :param message_id: Unique identifier for the message (default: generated UUID).
+        :param parent_id: ID of the user message that spawned child messages.
+        :param conversation_id: ID of the conversation this message belongs to (default: None).
+        :param content: Message content.
+        :param timestamp: Timestamp of the message (default: current UTC time).
+        :return: A WebSocketObservabilityTraceMessage instance.
+        """
+        try:
+            return WebSocketObservabilityTraceMessage(id=message_id,
+                                                      parent_id=parent_id,
+                                                      conversation_id=conversation_id,
+                                                      content=content,
+                                                      timestamp=timestamp)
+
+        except Exception as e:
+            logger.exception("Error creating observability trace message: %s", str(e))
+            return None

nat/front_ends/fastapi/register.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/fastapi/response_helpers.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -23,20 +23,20 @@ from nat.data_models.api_server import ResponseSerializable
 from nat.data_models.step_adaptor import StepAdaptorConfig
 from nat.front_ends.fastapi.intermediate_steps_subscriber import pull_intermediate
 from nat.front_ends.fastapi.step_adaptor import StepAdaptor
-from nat.runtime.session import SessionManager
+from nat.runtime.session import Session
 from nat.utils.producer_consumer_queue import AsyncIOProducerConsumerQueue
 
 
 async def generate_streaming_response_as_str(payload: typing.Any,
                                              *,
-                                             session_manager: SessionManager,
+                                             session: Session,
                                              streaming: bool,
                                              step_adaptor: StepAdaptor = StepAdaptor(StepAdaptorConfig()),
                                              result_type: type | None = None,
                                              output_type: type | None = None) -> AsyncGenerator[str]:
 
     async for item in generate_streaming_response(payload,
-                                                  session_manager=session_manager,
+                                                  session=session,
                                                   streaming=streaming,
                                                   step_adaptor=step_adaptor,
                                                   result_type=result_type,
@@ -51,13 +51,13 @@ async def generate_streaming_response_as_str(payload: typing.Any,
 
 async def generate_streaming_response(payload: typing.Any,
                                       *,
-                                      session_manager: SessionManager,
+                                      session: Session,
                                       streaming: bool,
                                       step_adaptor: StepAdaptor = StepAdaptor(StepAdaptorConfig()),
                                       result_type: type | None = None,
                                       output_type: type | None = None) -> AsyncGenerator[ResponseSerializable]:
 
-    async with session_manager.run(payload) as runner:
+    async with session.run(payload) as runner:
 
         q: AsyncIOProducerConsumerQueue[ResponseSerializable] = AsyncIOProducerConsumerQueue()
 
@@ -65,7 +65,7 @@ async def generate_streaming_response(payload: typing.Any,
         intermediate_complete = await pull_intermediate(q, step_adaptor)
 
         async def pull_result():
-            if session_manager.workflow.has_streaming_output and streaming:
+            if session.workflow.has_streaming_output and streaming:
                 async for chunk in runner.result_stream(to_type=output_type):
                     await q.put(chunk)
             else:
@@ -107,19 +107,20 @@ async def generate_streaming_response(payload: typing.Any,
 
 async def generate_single_response(
     payload: typing.Any,
-    session_manager: SessionManager,
+    session: Session,
     result_type: type | None = None,
 ) -> typing.Any:
-    if (not session_manager.workflow.has_single_output):
+
+    if not session.workflow.has_single_output:
         raise ValueError("Cannot get a single output value for streaming workflows")
 
-    async with session_manager.run(payload) as runner:
+    async with session.run(payload) as runner:
         return await runner.result(to_type=result_type)
 
 
 async def generate_streaming_response_full(payload: typing.Any,
                                            *,
-                                           session_manager: SessionManager,
+                                           session: Session,
                                            streaming: bool,
                                            result_type: type | None = None,
                                            output_type: type | None = None,
@@ -137,14 +138,14 @@ async def generate_streaming_response_full(payload: typing.Any,
         else:
             allowed_types = set(filter_steps.split(','))
 
-    async with session_manager.run(payload) as runner:
+    async with session.run(payload) as runner:
         q: AsyncIOProducerConsumerQueue[ResponseSerializable] = AsyncIOProducerConsumerQueue()
 
         # Start the intermediate stream without step adaptor
         intermediate_complete = await pull_intermediate(q, None)
 
         async def pull_result():
-            if session_manager.workflow.has_streaming_output and streaming:
+            if session.workflow.has_streaming_output and streaming:
                 async for chunk in runner.result_stream(to_type=output_type):
                     await q.put(chunk)
             else:
@@ -174,7 +175,7 @@ async def generate_streaming_response_full(payload: typing.Any,
 
 async def generate_streaming_response_full_as_str(payload: typing.Any,
                                                   *,
-                                                  session_manager: SessionManager,
+                                                  session: Session,
                                                   streaming: bool,
                                                   result_type: type | None = None,
                                                   output_type: type | None = None,
@@ -183,7 +184,7 @@ async def generate_streaming_response_full_as_str(payload: typing.Any,
     Similar to generate_streaming_response but converts the response to a string format.
     """
     async for item in generate_streaming_response_full(payload,
-                                                       session_manager=session_manager,
+                                                       session=session,
                                                        streaming=streaming,
                                                        result_type=result_type,
                                                        output_type=output_type,

nat/front_ends/fastapi/step_adaptor.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/fastapi/utils.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/register.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,4 +18,3 @@
 
 from .console import register as console_register
 from .fastapi import register as fastapi_register
-from .mcp import register as mcp_register

nat/front_ends/simple_base/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/simple_base/simple_front_end_plugin_base.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -47,9 +47,11 @@ class SimpleFrontEndPluginBase(FrontEndBase[FrontEndConfigT], ABC):
 
                 click.echo(stream.getvalue())
 
-            workflow = await builder.build()
-            session_manager = SessionManager(workflow)
-            await self.run_workflow(session_manager)
+            session_manager = await SessionManager.create(config=self.full_config, shared_builder=builder)
+            try:
+                await self.run_workflow(session_manager)
+            finally:
+                await session_manager.shutdown()
 
     @abstractmethod
     async def run_workflow(self, session_manager: SessionManager):

nat/llm/aws_bedrock_llm.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/llm/azure_openai_llm.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,6 +16,7 @@
 from pydantic import AliasChoices
 from pydantic import ConfigDict
 from pydantic import Field
+from pydantic import computed_field
 
 from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
@@ -59,6 +60,14 @@ class AzureOpenAIModelConfig(
                                            description="Top-p for distribution sampling.",
                                            space=SearchSpace(high=1.0, low=0.5, step=0.1))
 
+    @computed_field
+    @property
+    def model_name(self) -> str:
+        """
+        Returns the model name for compatibility with other parts of the code base which expect a model_name attribute.
+        """
+        return self.azure_deployment
+
 
 @register_llm_provider(config_type=AzureOpenAIModelConfig)
 async def azure_openai_llm(config: AzureOpenAIModelConfig, _builder: Builder):