nvidia-nat 1.3a20250819__py3-none-any.whl → 1.3.0a20250823__py3-none-any.whl

nat/eval/evaluate.py

@@ -42,7 +42,7 @@ from nat.runtime.session import SessionManager
 logger = logging.getLogger(__name__)
 
 
-class EvaluationRun:  # pylint: disable=too-many-public-methods
+class EvaluationRun:
     """
     Instantiated for each evaluation run and used to store data for that single run.
 
@@ -63,7 +63,16 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
 
         # Helpers
         self.intermediate_step_adapter: IntermediateStepAdapter = IntermediateStepAdapter()
-        self.weave_eval: WeaveEvaluationIntegration = WeaveEvaluationIntegration()
+
+        # Create evaluation trace context
+        try:
+            from nat.eval.utils.eval_trace_ctx import WeaveEvalTraceContext
+            self.eval_trace_context = WeaveEvalTraceContext()
+        except Exception:
+            from nat.eval.utils.eval_trace_ctx import EvalTraceContext
+            self.eval_trace_context = EvalTraceContext()
+
+        self.weave_eval: WeaveEvaluationIntegration = WeaveEvaluationIntegration(self.eval_trace_context)
         # Metadata
         self.eval_input: EvalInput | None = None
         self.workflow_interrupted: bool = False
@@ -310,7 +319,7 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
             except Exception as e:
                 logger.exception("Failed to delete old job directory: %s: %s", dir_to_delete, e, exc_info=True)
 
-    def write_output(self, dataset_handler: DatasetHandler, profiler_results: ProfilerResults):  # pylint: disable=unused-argument  # noqa: E501
+    def write_output(self, dataset_handler: DatasetHandler, profiler_results: ProfilerResults):
         workflow_output_file = self.eval_config.general.output_dir / "workflow_output.json"
         workflow_output_file.parent.mkdir(parents=True, exist_ok=True)
 
@@ -401,6 +410,33 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
 
         return workflow_type
 
+    async def wait_for_all_export_tasks_local(self, session_manager: SessionManager, timeout: float) -> None:
+        """Wait for all trace export tasks to complete for local workflows.
+
+        This only works for local workflows where we have direct access to the
+        SessionManager and its underlying workflow with exporter manager.
+        """
+        try:
+            workflow = session_manager.workflow
+            all_exporters = await workflow.get_all_exporters()
+            if not all_exporters:
+                logger.debug("No exporters to wait for")
+                return
+
+            logger.info("Waiting for export tasks from %d local exporters (timeout: %ds)", len(all_exporters), timeout)
+
+            for name, exporter in all_exporters.items():
+                try:
+                    await exporter.wait_for_tasks(timeout=timeout)
+                    logger.info("Export tasks completed for exporter: %s", name)
+                except Exception as e:
+                    logger.warning("Error waiting for export tasks from %s: %s", name, e)
+
+            logger.info("All local export task waiting completed")
+
+        except Exception as e:
+            logger.warning("Failed to wait for local export tasks: %s", e)
+
     async def run_and_evaluate(self,
                                session_manager: SessionManager | None = None,
                                job_id: str | None = None) -> EvaluationRunOutput:
@@ -442,11 +478,13 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
         dataset_config = self.eval_config.general.dataset  # Currently only one dataset is supported
         if not dataset_config:
             logger.info("No dataset found, nothing to evaluate")
-            return EvaluationRunOutput(
-                workflow_output_file=self.workflow_output_file,
-                evaluator_output_files=self.evaluator_output_files,
-                workflow_interrupted=self.workflow_interrupted,
-            )
+            return EvaluationRunOutput(workflow_output_file=self.workflow_output_file,
+                                       evaluator_output_files=self.evaluator_output_files,
+                                       workflow_interrupted=self.workflow_interrupted,
+                                       eval_input=EvalInput(eval_input_items=[]),
+                                       evaluation_results=[],
+                                       usage_stats=UsageStats(),
+                                       profiler_results=ProfilerResults())
 
         dataset_handler = DatasetHandler(dataset_config=dataset_config,
                                          reps=self.config.reps,
@@ -456,11 +494,13 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
         self.eval_input = dataset_handler.get_eval_input_from_dataset(self.config.dataset)
         if not self.eval_input.eval_input_items:
             logger.info("Dataset is empty. Nothing to evaluate.")
-            return EvaluationRunOutput(
-                workflow_output_file=self.workflow_output_file,
-                evaluator_output_files=self.evaluator_output_files,
-                workflow_interrupted=self.workflow_interrupted,
-            )
+            return EvaluationRunOutput(workflow_output_file=self.workflow_output_file,
+                                       evaluator_output_files=self.evaluator_output_files,
+                                       workflow_interrupted=self.workflow_interrupted,
+                                       eval_input=self.eval_input,
+                                       evaluation_results=self.evaluation_results,
+                                       usage_stats=self.usage_stats,
+                                       profiler_results=ProfilerResults())
 
         # Run workflow and evaluate
         async with WorkflowEvalBuilder.from_config(config=config) as eval_workflow:
@@ -468,18 +508,22 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
             self.weave_eval.initialize_logger(workflow_alias, self.eval_input, config)
 
             # Run workflow
-            if self.config.endpoint:
-                await self.run_workflow_remote()
-            else:
-                if not self.config.skip_workflow:
+            with self.eval_trace_context.evaluation_context():
+                if self.config.endpoint:
+                    await self.run_workflow_remote()
+                elif not self.config.skip_workflow:
                     if session_manager is None:
                         session_manager = SessionManager(eval_workflow.build(),
                                                          max_concurrency=self.eval_config.general.max_concurrency)
                     await self.run_workflow_local(session_manager)
 
-            # Evaluate
-            evaluators = {name: eval_workflow.get_evaluator(name) for name in self.eval_config.evaluators}
-            await self.run_evaluators(evaluators)
+                # Evaluate
+                evaluators = {name: eval_workflow.get_evaluator(name) for name in self.eval_config.evaluators}
+                await self.run_evaluators(evaluators)
+
+                # Wait for all trace export tasks to complete (local workflows only)
+                if session_manager and not self.config.endpoint:
+                    await self.wait_for_all_export_tasks_local(session_manager, timeout=self.config.export_timeout)
 
         # Profile the workflow
         profiler_results = await self.profile_workflow()

nat/eval/rag_evaluator/register.py

@@ -73,7 +73,7 @@ class RagasEvaluatorConfig(EvaluatorBaseConfig, name="ragas"):
         if isinstance(self.metric, str):
             return self.metric
         if isinstance(self.metric, dict) and self.metric:
-            return next(iter(self.metric.keys()))  # pylint: disable=no-member
+            return next(iter(self.metric.keys()))
         return ""
 
     @property
@@ -82,7 +82,7 @@ class RagasEvaluatorConfig(EvaluatorBaseConfig, name="ragas"):
         if isinstance(self.metric, str):
             return RagasMetricConfig()  # Default config when only a metric name is given
         if isinstance(self.metric, dict) and self.metric:
-            return next(iter(self.metric.values()))  # pylint: disable=no-member
+            return next(iter(self.metric.values()))
         return RagasMetricConfig()  # Default config when an invalid type is provided
 
 

nat/eval/register.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 # flake8: noqa
-# pylint: disable=unused-import
 
 # Import evaluators which need to be automatically registered here
 from .rag_evaluator.register import register_ragas_evaluator

nat/eval/tunable_rag_evaluator/evaluate.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
 import logging
 from typing import Callable
 
@@ -23,7 +22,6 @@ from langchain.schema import HumanMessage
 from langchain.schema import SystemMessage
 from langchain_core.language_models import BaseChatModel
 from langchain_core.runnables import RunnableLambda
-from tqdm import tqdm
 
 from nat.eval.evaluator.base_evaluator import BaseEvaluator
 from nat.eval.evaluator.evaluator_model import EvalInputItem
@@ -31,7 +29,6 @@ from nat.eval.evaluator.evaluator_model import EvalOutputItem
 
 logger = logging.getLogger(__name__)
 
-# pylint: disable=line-too-long
 # flake8: noqa: E501
 
 

nat/eval/utils/eval_trace_ctx.py

@@ -0,0 +1,89 @@
+# 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 logging
+from collections.abc import Callable
+from contextlib import contextmanager
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+# Type alias for evaluation call objects that have an optional 'id' attribute
+EvalCallType = Any  # Could be Weave Call object or other tracing framework objects
+
+
+class EvalTraceContext:
+    """
+    Evaluation trace context manager for coordinating traces.
+
+    This class provides a framework-agnostic way to:
+    1. Track evaluation calls/contexts
+    2. Ensure proper parent-child relationships in traces
+    """
+
+    def __init__(self):
+        self.eval_call: EvalCallType | None = None  # Store the evaluation call/context for propagation
+
+    def set_eval_call(self, eval_call: EvalCallType | None) -> None:
+        """Set the evaluation call/context for propagation to traces."""
+        self.eval_call = eval_call
+        if eval_call:
+            logger.debug("Set evaluation call context: %s", getattr(eval_call, 'id', str(eval_call)))
+
+    def get_eval_call(self) -> EvalCallType | None:
+        """Get the current evaluation call/context."""
+        return self.eval_call
+
+    @contextmanager
+    def evaluation_context(self):
+        """
+        Context manager that can be overridden by framework-specific implementations.
+        Default implementation is a no-op.
+        """
+        yield
+
+
+class WeaveEvalTraceContext(EvalTraceContext):
+    """
+    Weave-specific implementation of evaluation trace context.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.available = False
+        self.set_call_stack: Callable[[list[EvalCallType]], Any] | None = None
+
+        try:
+            from weave.trace.context.call_context import set_call_stack
+            self.set_call_stack = set_call_stack
+            self.available = True
+        except ImportError:
+            self.available = False
+            logger.debug("Weave not available for trace context")
+
+    @contextmanager
+    def evaluation_context(self):
+        """Set the evaluation call as active context for Weave traces."""
+        if self.available and self.eval_call and self.set_call_stack:
+            try:
+                with self.set_call_stack([self.eval_call]):
+                    logger.debug("Set Weave evaluation call context: %s",
+                                 getattr(self.eval_call, 'id', str(self.eval_call)))
+                    yield
+            except Exception as e:
+                logger.warning("Failed to set Weave evaluation call context: %s", e)
+                yield
+        else:
+            yield

nat/eval/utils/weave_eval.py

@@ -15,6 +15,7 @@
 
 import asyncio
 import logging
+from typing import TYPE_CHECKING
 from typing import Any
 
 from nat.eval.evaluator.evaluator_model import EvalInput
@@ -24,26 +25,30 @@ from nat.eval.usage_stats import UsageStats
 from nat.eval.usage_stats import UsageStatsItem
 from nat.profiler.data_models import ProfilerResults
 
+if TYPE_CHECKING:
+    from nat.eval.utils.eval_trace_ctx import EvalTraceContext
+
 logger = logging.getLogger(__name__)
 
 
-class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
+class WeaveEvaluationIntegration:
     """
     Class to handle all Weave integration functionality.
     """
 
-    def __init__(self):
+    def __init__(self, eval_trace_context: "EvalTraceContext"):
         self.available = False
         self.client = None
         self.eval_logger = None
         self.pred_loggers = {}
+        self.eval_trace_context = eval_trace_context
 
         try:
             from weave.flow.eval_imperative import EvaluationLogger
             from weave.flow.eval_imperative import ScoreLogger
             from weave.trace.context import weave_client_context
-            self.EvaluationLogger = EvaluationLogger
-            self.ScoreLogger = ScoreLogger
+            self.evaluation_logger_cls = EvaluationLogger
+            self.score_logger_cls = ScoreLogger
             self.weave_client_context = weave_client_context
             self.available = True
         except ImportError:
@@ -89,9 +94,12 @@ class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
             weave_dataset = self._get_weave_dataset(eval_input)
             config_dict = config.model_dump(mode="json")
             config_dict["name"] = workflow_alias
-            self.eval_logger = self.EvaluationLogger(model=config_dict, dataset=weave_dataset)
+            self.eval_logger = self.evaluation_logger_cls(model=config_dict, dataset=weave_dataset)
             self.pred_loggers = {}
 
+            # Capture the current evaluation call for context propagation
+            self.eval_trace_context.set_eval_call(self.eval_logger._evaluate_call)
+
             return True
         except Exception as e:
             self.eval_logger = None
@@ -137,7 +145,7 @@ class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
             await asyncio.gather(*coros)
 
     async def afinish_loggers(self):
-        """Finish all prediction loggers."""
+        """Finish all prediction loggers and wait for exports."""
         if not self.eval_logger:
             return
 
@@ -157,7 +165,6 @@ class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
         if profiler_results.workflow_runtime_metrics:
             profile_metrics["wf_runtime_p95"] = profiler_results.workflow_runtime_metrics.p95
 
-        # TODO:get the LLM tokens from the usage stats and log them
         profile_metrics["total_runtime"] = usage_stats.total_runtime
 
         return profile_metrics

nat/experimental/test_time_compute/models/strategy_base.py

@@ -17,9 +17,10 @@ from abc import ABC
 from abc import abstractmethod
 
 from nat.builder.builder import Builder
-from nat.experimental.test_time_compute.models.ttc_item import TTCItem
-from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum, PipelineTypeEnum
 from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
 
 
 class StrategyBase(ABC):

nat/experimental/test_time_compute/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 
 from .editing import iterative_plan_refinement_editor

nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py

@@ -135,8 +135,6 @@ class LLMBasedOutputMergingSelector(StrategyBase):
             except Exception as e:
                 logger.error(f"Error parsing merged output: {e}")
                 raise ValueError("Failed to parse merged output.")
-        else:
-            merged_output = merged_output
 
         logger.info("Merged output: %s", str(merged_output))
 

nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py

@@ -307,7 +307,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
         async def start_evaluation(request: EvaluateRequest, background_tasks: BackgroundTasks, http_request: Request):
             """Handle evaluation requests."""
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 # if job_id is present and already exists return the job info
                 if request.job_id:
@@ -336,7 +336,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
             """Get the status of an evaluation job."""
             logger.info("Getting status for job %s", job_id)
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 job = job_store.get_job(job_id)
                 if not job:
@@ -349,7 +349,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
             """Get the status of the last created evaluation job."""
             logger.info("Getting last job status")
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 job = job_store.get_last_job()
                 if not job:
@@ -361,7 +361,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
         async def get_jobs(http_request: Request, status: str | None = None) -> list[EvaluateStatusResponse]:
             """Get all jobs, optionally filtered by status."""
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 if status is None:
                     logger.info("Getting all jobs")
@@ -522,9 +522,9 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
         workflow = session_manager.workflow
 
-        GenerateBodyType = workflow.input_schema  # pylint: disable=invalid-name
-        GenerateStreamResponseType = workflow.streaming_output_schema  # pylint: disable=invalid-name
-        GenerateSingleResponseType = workflow.single_output_schema  # pylint: disable=invalid-name
+        GenerateBodyType = workflow.input_schema
+        GenerateStreamResponseType = workflow.streaming_output_schema
+        GenerateSingleResponseType = workflow.single_output_schema
 
         # Append job_id and expiry_seconds to the input schema, this effectively makes these reserved keywords
         # Consider prefixing these with "nat_" to avoid conflicts
@@ -572,7 +572,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
                 response.headers["Content-Type"] = "application/json"
 
-                async with session_manager.session(request=request,
+                async with session_manager.session(http_connection=request,
                                                    user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return await generate_single_response(None, session_manager, result_type=result_type)
@@ -583,7 +583,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def get_stream(request: Request):
 
-                async with session_manager.session(request=request,
+                async with session_manager.session(http_connection=request,
                                                    user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
@@ -618,7 +618,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
                 response.headers["Content-Type"] = "application/json"
 
-                async with session_manager.session(request=request,
+                async with session_manager.session(http_connection=request,
                                                    user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return await generate_single_response(payload, session_manager, result_type=result_type)
@@ -632,7 +632,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def post_stream(request: Request, payload: request_type):
 
-                async with session_manager.session(request=request,
+                async with session_manager.session(http_connection=request,
                                                    user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
@@ -677,7 +677,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                 # Check if streaming is requested
                 stream_requested = getattr(payload, 'stream', False)
 
-                async with session_manager.session(request=request):
+                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"},
@@ -688,42 +688,41 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                                                      step_adaptor=self.get_step_adaptor(),
                                                      result_type=ChatResponseChunk,
                                                      output_type=ChatResponseChunk))
-                    else:
-                        # Return single response - check if workflow supports non-streaming
-                        try:
+
+                    # 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 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
-                            else:
-                                raise
+                            return single_response
+                        raise
 
             return post_openai_api_compatible
 
@@ -758,7 +757,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                     http_request: Request) -> AsyncGenerateResponse | AsyncGenerationStatusResponse:
                 """Handle async generation requests."""
 
-                async with session_manager.session(request=http_request):
+                async with session_manager.session(http_connection=http_request):
 
                     # if job_id is present and already exists return the job info
                     if request.job_id:
@@ -804,7 +803,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
             """Get the status of an async job."""
             logger.info("Getting status for job %s", job_id)
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 job = job_store.get_job(job_id)
                 if not job:

nat/front_ends/fastapi/message_handler.py

@@ -86,7 +86,7 @@ class WebSocketMessageHandler:
 
     async def __aexit__(self, exc_type, exc_value, traceback) -> None:
 
-        # TODO: Handle the exit  # pylint: disable=fixme
+        # TODO: Handle the exit
         pass
 
     async def run(self) -> None:
@@ -105,12 +105,10 @@ class WebSocketMessageHandler:
                 if (isinstance(validated_message, WebSocketUserMessage)):
                     await self.process_workflow_request(validated_message)
 
-                elif isinstance(
-                        validated_message,
-                    (  # noqa: E131
-                        WebSocketSystemResponseTokenMessage,
-                        WebSocketSystemIntermediateStepMessage,
-                        WebSocketSystemInteractionMessage)):
+                elif isinstance(validated_message,
+                                (WebSocketSystemResponseTokenMessage,
+                                 WebSocketSystemIntermediateStepMessage,
+                                 WebSocketSystemInteractionMessage)):
                     # These messages are already handled by self.create_websocket_message(data_model=value, …)
                     # No further processing is needed here.
                     pass
@@ -119,11 +117,9 @@ class WebSocketMessageHandler:
                     user_content = await self.process_user_message_content(validated_message)
                     self._user_interaction_response.set_result(user_content)
             except (asyncio.CancelledError, WebSocketDisconnect):
-                # TODO: Handle the disconnect  # pylint: disable=fixme
+                # TODO: Handle the disconnect
                 break
 
-        return None
-
     async def process_user_message_content(
             self, user_content: WebSocketUserMessage | WebSocketUserInteractionResponseMessage) -> BaseModel | None:
         """
@@ -162,12 +158,13 @@ class WebSocketMessageHandler:
 
             if isinstance(content, TextContent) and (self._running_workflow_task is None):
 
-                def _done_callback(task: asyncio.Task):  # pylint: disable=unused-argument
+                def _done_callback(task: asyncio.Task):
                     self._running_workflow_task = None
 
                 self._running_workflow_task = asyncio.create_task(
-                    self._run_workflow(content.text,
-                                       self._conversation_id,
+                    self._run_workflow(payload=content.text,
+                                       user_message_id=self._message_parent_id,
+                                       conversation_id=self._conversation_id,
                                        result_type=self._schema_output_mapping[self._workflow_schema_type],
                                        output_type=self._schema_output_mapping[
                                            self._workflow_schema_type])).add_done_callback(_done_callback)
@@ -290,14 +287,16 @@ class WebSocketMessageHandler:
 
     async def _run_workflow(self,
                             payload: typing.Any,
+                            user_message_id: str | None = None,
                             conversation_id: str | None = None,
                             result_type: type | None = None,
                             output_type: type | None = None) -> None:
 
         try:
             async with self._session_manager.session(
+                    user_message_id=user_message_id,
                     conversation_id=conversation_id,
-                    request=self._socket,
+                    http_connection=self._socket,
                     user_input_callback=self.human_interaction_callback,
                     user_authentication_callback=(self._flow_handler.authenticate
                                                   if self._flow_handler else None)) as session: