nvidia-nat 1.2.0a20250813__py3-none-any.whl

nat/front_ends/fastapi/message_handler.py

@@ -0,0 +1,309 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+import typing
+import uuid
+from typing import Any
+
+from fastapi import WebSocket
+from pydantic import BaseModel
+from pydantic import ValidationError
+from starlette.websockets import WebSocketDisconnect
+
+from nat.authentication.interfaces import FlowHandlerBase
+from nat.data_models.api_server import ChatResponse
+from nat.data_models.api_server import Error
+from nat.data_models.api_server import ErrorTypes
+from nat.data_models.api_server import ResponsePayloadOutput
+from nat.data_models.api_server import ResponseSerializable
+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 WebSocketSystemInteractionMessage
+from nat.data_models.api_server import WebSocketSystemIntermediateStepMessage
+from nat.data_models.api_server import WebSocketSystemResponseTokenMessage
+from nat.data_models.api_server import WebSocketUserInteractionResponseMessage
+from nat.data_models.api_server import WebSocketUserMessage
+from nat.data_models.interactive import HumanPromptNotification
+from nat.data_models.interactive import HumanResponse
+from nat.data_models.interactive import HumanResponseNotification
+from nat.data_models.interactive import InteractionPrompt
+from nat.front_ends.fastapi.message_validator import MessageValidator
+from nat.front_ends.fastapi.response_helpers import generate_streaming_response
+from nat.front_ends.fastapi.step_adaptor import StepAdaptor
+from nat.runtime.session import SessionManager
+
+logger = logging.getLogger(__name__)
+
+
+class WebSocketMessageHandler:
+
+    def __init__(self, socket: WebSocket, session_manager: SessionManager, step_adaptor: StepAdaptor):
+        self._socket: WebSocket = socket
+        self._session_manager: SessionManager = session_manager
+        self._step_adaptor: StepAdaptor = step_adaptor
+
+        self._message_validator: MessageValidator = MessageValidator()
+        self._running_workflow_task: asyncio.Task | None = None
+        self._message_parent_id: str = "default_id"
+        self._conversation_id: str | None = None
+        self._workflow_schema_type: str = None
+        self._user_interaction_response: asyncio.Future[HumanResponse] | None = None
+
+        self._flow_handler: FlowHandlerBase | None = None
+
+    def set_flow_handler(self, flow_handler: FlowHandlerBase) -> None:
+        self._flow_handler = flow_handler
+
+    async def __aenter__(self) -> "WebSocketMessageHandler":
+        await self._socket.accept()
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
+
+        # TODO: Handle the exit  # pylint: disable=fixme
+        pass
+
+    async def run(self) -> None:
+        """
+        Processes received messages from websocket and routes them appropriately.
+        """
+        while True:
+
+            try:
+
+                message: dict[str, Any] = await self._socket.receive_json()
+
+                validated_message: BaseModel = await self._message_validator.validate_message(message)
+
+                # Received a request to start a workflow
+                if (isinstance(validated_message, WebSocketUserMessage)):
+                    await self.process_workflow_request(validated_message)
+
+                elif isinstance(
+                        validated_message,
+                    (  # noqa: E131
+                        WebSocketSystemResponseTokenMessage,
+                        WebSocketSystemIntermediateStepMessage,
+                        WebSocketSystemInteractionMessage)):
+                    # These messages are already handled by self.create_websocket_message(data_model=value, …)
+                    # No further processing is needed here.
+                    pass
+
+                elif (isinstance(validated_message, WebSocketUserInteractionResponseMessage)):
+                    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
+                break
+
+        return None
+
+    async def process_user_message_content(
+            self, user_content: WebSocketUserMessage | WebSocketUserInteractionResponseMessage) -> BaseModel | None:
+        """
+        Processes the contents of a user message.
+
+        :param user_content: Incoming content data model.
+        :return: A validated Pydantic user content model or None if not found.
+        """
+
+        for user_message in user_content.content.messages[::-1]:
+            if (user_message.role == "user"):
+
+                for attachment in user_message.content:
+
+                    if isinstance(attachment, TextContent):
+                        return attachment
+
+        return None
+
+    async def process_workflow_request(self, user_message_as_validated_type: WebSocketUserMessage) -> None:
+        """
+        Process user messages and routes them appropriately.
+
+        :param user_message_as_validated_type: A WebSocketUserMessage Data Model instance.
+        """
+
+        try:
+            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
+
+            content: BaseModel | None = await self.process_user_message_content(user_message_as_validated_type)
+
+            if content is None:
+                raise ValueError(f"User message content could not be found: {user_message_as_validated_type}")
+
+            if isinstance(content, TextContent) and (self._running_workflow_task is None):
+
+                def _done_callback(task: asyncio.Task):  # pylint: disable=unused-argument
+                    self._running_workflow_task = None
+
+                # await self._process_response()
+                self._running_workflow_task = asyncio.create_task(
+                    self._run_workflow(content.text, self._conversation_id,
+                                       result_type=ChatResponse)).add_done_callback(_done_callback)
+
+        except ValueError as e:
+            logger.error("User message content not found: %s", str(e), exc_info=True)
+            await self.create_websocket_message(data_model=Error(code=ErrorTypes.INVALID_USER_MESSAGE_CONTENT,
+                                                                 message="User message content could not be found",
+                                                                 details=str(e)),
+                                                message_type=WebSocketMessageType.ERROR_MESSAGE,
+                                                status=WebSocketMessageStatus.IN_PROGRESS)
+
+    async def create_websocket_message(self,
+                                       data_model: BaseModel,
+                                       message_type: str | None = None,
+                                       status: str = WebSocketMessageStatus.IN_PROGRESS) -> None:
+        """
+        Creates a websocket message that will be ready for routing based on message type or data model.
+
+        :param data_model: Message content model.
+        :param message_type: Message content model.
+        :param status: Message content model.
+        """
+        try:
+            message: BaseModel | None = None
+
+            if message_type is None:
+                message_type = await self._message_validator.resolve_message_type_by_data(data_model)
+
+            message_schema: type[BaseModel] = await self._message_validator.get_message_schema_by_type(message_type)
+
+            if 'id' in data_model.model_fields:
+                message_id: str = data_model.id
+            else:
+                message_id = str(uuid.uuid4())
+
+            content: BaseModel = await self._message_validator.convert_data_to_message_content(data_model)
+
+            if issubclass(message_schema, WebSocketSystemResponseTokenMessage):
+                message = await self._message_validator.create_system_response_token_message(
+                    message_id=message_id,
+                    parent_id=self._message_parent_id,
+                    conversation_id=self._conversation_id,
+                    content=content,
+                    status=status)
+
+            elif issubclass(message_schema, WebSocketSystemIntermediateStepMessage):
+                message = await self._message_validator.create_system_intermediate_step_message(
+                    message_id=message_id,
+                    parent_id=await self._message_validator.get_intermediate_step_parent_id(data_model),
+                    conversation_id=self._conversation_id,
+                    content=content,
+                    status=status)
+
+            elif issubclass(message_schema, WebSocketSystemInteractionMessage):
+                message = await self._message_validator.create_system_interaction_message(
+                    message_id=message_id,
+                    parent_id=self._message_parent_id,
+                    conversation_id=self._conversation_id,
+                    content=content,
+                    status=status)
+
+            elif isinstance(content, Error):
+                raise ValidationError(f"Invalid input data creating websocket message. {data_model.model_dump_json()}")
+
+            elif issubclass(message_schema, Error):
+                raise TypeError(f"Invalid message type: {message_type}")
+
+            elif (message is None):
+                raise ValueError(
+                    f"Message type could not be resolved by input data model: {data_model.model_dump_json()}")
+
+        except (ValidationError, TypeError, ValueError) as e:
+            logger.error("A data vaidation error ocurred creating websocket message: %s", str(e), exc_info=True)
+            message = await self._message_validator.create_system_response_token_message(
+                message_type=WebSocketMessageType.ERROR_MESSAGE,
+                conversation_id=self._conversation_id,
+                content=Error(code=ErrorTypes.UNKNOWN_ERROR, message="default", details=str(e)))
+
+        finally:
+            if (message is not None):
+                await self._socket.send_json(message.model_dump())
+
+    async def human_interaction_callback(self, prompt: InteractionPrompt) -> HumanResponse:
+        """
+        Registered human interaction callback that processes human interactions and returns
+        responses from websocket connection.
+
+        :param prompt: Incoming interaction content data model.
+        :return: A Text Content Base Pydantic model.
+        """
+
+        # First create a future from the loop for the human response
+        human_response_future: asyncio.Future[HumanResponse] = asyncio.get_running_loop().create_future()
+
+        # Then add the future to the outstanding human prompts dictionary
+        self._user_interaction_response = human_response_future
+
+        try:
+
+            await self.create_websocket_message(data_model=prompt.content,
+                                                message_type=WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE,
+                                                status=WebSocketMessageStatus.IN_PROGRESS)
+
+            if (isinstance(prompt.content, HumanPromptNotification)):
+
+                return HumanResponseNotification()
+
+            # Wait for the human response future to complete
+            interaction_response: HumanResponse = await human_response_future
+
+            interaction_response: HumanResponse = await self._message_validator.convert_text_content_to_human_response(
+                interaction_response, prompt.content)
+
+            return interaction_response
+
+        finally:
+            # Delete the future from the outstanding human prompts dictionary
+            self._user_interaction_response = None
+
+    async def _run_workflow(self,
+                            payload: typing.Any,
+                            conversation_id: str | None = None,
+                            result_type: type | None = None,
+                            output_type: type | None = None) -> None:
+
+        try:
+            async with self._session_manager.session(
+                    conversation_id=conversation_id,
+                    request=self._socket,
+                    user_input_callback=self.human_interaction_callback,
+                    user_authentication_callback=(self._flow_handler.authenticate
+                                                  if self._flow_handler else None)) as session:
+
+                async for value in generate_streaming_response(payload,
+                                                               session_manager=session,
+                                                               streaming=True,
+                                                               step_adaptor=self._step_adaptor,
+                                                               result_type=result_type,
+                                                               output_type=output_type):
+
+                    if not isinstance(value, ResponseSerializable):
+                        value = ResponsePayloadOutput(payload=value)
+
+                    await self.create_websocket_message(data_model=value, status=WebSocketMessageStatus.IN_PROGRESS)
+
+        finally:
+            await self.create_websocket_message(data_model=SystemResponseContent(),
+                                                message_type=WebSocketMessageType.RESPONSE_MESSAGE,
+                                                status=WebSocketMessageStatus.COMPLETE)

nat/front_ends/fastapi/message_validator.py

@@ -0,0 +1,354 @@
+# 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 datetime
+import json
+import logging
+import uuid
+from typing import Any
+from typing import Literal
+
+from pydantic import BaseModel
+from pydantic import ValidationError
+
+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 ResponseIntermediateStep
+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 WebSocketSystemInteractionMessage
+from nat.data_models.api_server import WebSocketSystemIntermediateStepMessage
+from nat.data_models.api_server import WebSocketSystemResponseTokenMessage
+from nat.data_models.api_server import WebSocketUserInteractionResponseMessage
+from nat.data_models.api_server import WebSocketUserMessage
+from nat.data_models.api_server import WorkflowSchemaType
+from nat.data_models.interactive import BinaryHumanPromptOption
+from nat.data_models.interactive import HumanPrompt
+from nat.data_models.interactive import HumanPromptBase
+from nat.data_models.interactive import HumanPromptBinary
+from nat.data_models.interactive import HumanPromptCheckbox
+from nat.data_models.interactive import HumanPromptDropdown
+from nat.data_models.interactive import HumanPromptRadio
+from nat.data_models.interactive import HumanPromptText
+from nat.data_models.interactive import HumanResponse
+from nat.data_models.interactive import HumanResponseBinary
+from nat.data_models.interactive import HumanResponseCheckbox
+from nat.data_models.interactive import HumanResponseDropdown
+from nat.data_models.interactive import HumanResponseRadio
+from nat.data_models.interactive import HumanResponseText
+from nat.data_models.interactive import MultipleChoiceOption
+
+logger = logging.getLogger(__name__)
+
+
+class MessageValidator:
+
+    def __init__(self):
+        self._message_type_schema_mapping: dict[str, type[BaseModel]] = {
+            WebSocketMessageType.USER_MESSAGE: WebSocketUserMessage,
+            WebSocketMessageType.RESPONSE_MESSAGE: WebSocketSystemResponseTokenMessage,
+            WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE: WebSocketSystemIntermediateStepMessage,
+            WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE: WebSocketSystemInteractionMessage,
+            WebSocketMessageType.USER_INTERACTION_MESSAGE: WebSocketUserInteractionResponseMessage,
+            WebSocketMessageType.ERROR_MESSAGE: Error
+        }
+        self._data_type_schema_mapping: dict[str, type[BaseModel]] = {
+            WorkflowSchemaType.GENERATE: ResponsePayloadOutput,
+            WorkflowSchemaType.CHAT: ChatResponse,
+            WorkflowSchemaType.CHAT_STREAM: ChatResponseChunk,
+            WorkflowSchemaType.GENERATE_STREAM: ResponseIntermediateStep,
+        }
+        self._message_parent_id: str = "default_id"
+
+    async def validate_message(self, message: dict[str, Any]) -> BaseModel:
+        """
+        Validates an incoming WebSocket message against its expected schema.
+        If validation fails, returns a system response error message.
+
+        :param message: Incoming WebSocket message as a dictionary.
+        :return: A validated Pydantic model.
+        """
+        validated_message: BaseModel
+
+        try:
+            message_type = message.get("type")
+            if not message_type:
+                raise ValueError(f"Missing message type: {json.dumps(message)}")
+
+            schema: type[BaseModel] = await self.get_message_schema_by_type(message_type)
+
+            if issubclass(schema, Error):
+                raise TypeError(
+                    f"An error was encountered processing an incoming WebSocket message of type: {message_type}")
+
+            validated_message = schema(**message)
+            return validated_message
+
+        except (ValidationError, TypeError, ValueError) as e:
+            logger.error("A data validation error %s occurred for message: %s", str(e), str(message), exc_info=True)
+            return await self.create_system_response_token_message(message_type=WebSocketMessageType.ERROR_MESSAGE,
+                                                                   content=Error(code=ErrorTypes.INVALID_MESSAGE,
+                                                                                 message="Error validating message.",
+                                                                                 details=str(e)))
+
+    async def get_message_schema_by_type(self, message_type: str) -> type[BaseModel]:
+        """
+        Retrieves the corresponding Pydantic model schema based on the message type.
+
+        :param message_type: The type of message as a string.
+        :return: A Pydantic schema class if found, otherwise None.
+        """
+        try:
+            schema: type[BaseModel] | None = self._message_type_schema_mapping.get(message_type)
+
+            if schema is None:
+                raise ValueError(f"Unknown message type: {message_type}")
+
+            return schema
+
+        except (TypeError, ValueError) as e:
+            logger.error("Error retrieving schema for message type '%s': %s", message_type, str(e), exc_info=True)
+            return Error
+
+    async def convert_data_to_message_content(self, data_model: BaseModel) -> BaseModel:
+        """
+        Converts a Pydantic data model to a WebSocket message content instance.
+
+        :param data_model: Pydantic Data Model instance.
+        :return: A WebSocket Message Content Data Model instance.
+        """
+
+        validated_message_content: BaseModel = None
+        try:
+            if (isinstance(data_model, ResponsePayloadOutput)):
+                validated_message_content = SystemResponseContent(text=data_model.payload)
+
+            elif (isinstance(data_model, (ChatResponse, ChatResponseChunk))):
+                validated_message_content = SystemResponseContent(text=data_model.choices[0].message.content)
+
+            elif (isinstance(data_model, ResponseIntermediateStep)):
+                validated_message_content = SystemIntermediateStepContent(name=data_model.name,
+                                                                          payload=data_model.payload)
+            elif (isinstance(data_model, HumanPromptBase)):
+                validated_message_content = data_model
+            elif (isinstance(data_model, SystemResponseContent)):
+                return data_model
+            else:
+                raise ValueError(
+                    f"Input data could not be converted to validated message content: {data_model.model_dump_json()}")
+
+            return validated_message_content
+
+        except ValueError as e:
+            logger.error("Input data could not be converted to validated message content: %s", str(e), exc_info=True)
+            return Error(code=ErrorTypes.INVALID_DATA_CONTENT, message="Input data not supported.", details=str(e))
+
+    async def convert_text_content_to_human_response(self, text_content: TextContent,
+                                                     human_prompt: HumanPromptBase) -> HumanResponse:
+        """
+        Converts Message Text Content data model to a Human Response Base data model instance.
+
+        :param text_content: Pydantic TextContent Data Model instance.
+        :param human_prompt: Pydantic HumanPrompt Data Model instance.
+        :return: A Human Response Data Model instance.
+        """
+
+        human_response: HumanResponse = None
+        try:
+            if (isinstance(human_prompt, HumanPromptText)):
+                human_response = HumanResponseText(text=text_content.text)
+
+            elif (isinstance(human_prompt, HumanPromptBinary)):
+                human_response = HumanResponseBinary(selected_option=BinaryHumanPromptOption(value=text_content.text))
+
+            elif (isinstance(human_prompt, HumanPromptRadio)):
+                human_response = HumanResponseRadio(selected_option=MultipleChoiceOption(value=text_content.text))
+
+            elif (isinstance(human_prompt, HumanPromptCheckbox)):
+                human_response = HumanResponseCheckbox(selected_option=MultipleChoiceOption(value=text_content.text))
+
+            elif (isinstance(human_prompt, HumanPromptDropdown)):
+                human_response = HumanResponseDropdown(selected_option=MultipleChoiceOption(value=text_content.text))
+            else:
+                raise ValueError("Message content type not found")
+
+            return human_response
+
+        except ValueError as e:
+            logger.error("Error human response content not found: %s", str(e), exc_info=True)
+            return HumanResponseText(text=str(e))
+
+    async def resolve_message_type_by_data(self, data_model: BaseModel) -> str:
+        """
+        Resolve message type from a validated model
+
+        :param data_model: Pydantic Data Model instance.
+        :return: A WebSocket Message Content Data Model instance.
+        """
+
+        validated_message_type: str = ""
+        try:
+            if (isinstance(data_model, (ResponsePayloadOutput, ChatResponse, ChatResponseChunk))):
+                validated_message_type = WebSocketMessageType.RESPONSE_MESSAGE
+
+            elif (isinstance(data_model, ResponseIntermediateStep)):
+                validated_message_type = WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE
+
+            elif (isinstance(data_model, HumanPromptBase)):
+                validated_message_type = WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE
+            else:
+                raise ValueError("Data type not found")
+
+            return validated_message_type
+
+        except ValueError as e:
+            logger.error("Error type not found converting data to validated websocket message content: %s",
+                         str(e),
+                         exc_info=True)
+            return WebSocketMessageType.ERROR_MESSAGE
+
+    async def get_intermediate_step_parent_id(self, data_model: ResponseIntermediateStep) -> str:
+        """
+        Retrieves intermediate step parent_id from ResponseIntermediateStep instance.
+
+        :param data_model: ResponseIntermediateStep Data Model instance.
+        :return: Intermediate step parent_id or "default".
+        """
+        return data_model.parent_id or "root"
+
+    async def create_system_response_token_message(  # pylint: disable=R0917:too-many-positional-arguments
+        self,
+        message_type: Literal[WebSocketMessageType.RESPONSE_MESSAGE,
+                              WebSocketMessageType.ERROR_MESSAGE] = WebSocketMessageType.RESPONSE_MESSAGE,
+        message_id: str | None = str(uuid.uuid4()),
+        thread_id: str = "default",
+        parent_id: str = "default",
+        conversation_id: str | None = None,
+        content: SystemResponseContent
+        | Error = SystemResponseContent(),
+        status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
+        timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    ) -> WebSocketSystemResponseTokenMessage | None:
+        """
+        Creates a system response token message with default values.
+
+        :param message_type: Type of WebSocket message.
+        :param message_id: Unique identifier for the message (default: generated UUID).
+        :param thread_id: ID of the thread the message belongs to (default: "default").
+        :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 status: Status of the message (default: IN_PROGRESS).
+        :param timestamp: Timestamp of the message (default: current UTC time).
+        :return: A WebSocketSystemResponseTokenMessage instance.
+        """
+        try:
+            return WebSocketSystemResponseTokenMessage(type=message_type,
+                                                       id=message_id,
+                                                       thread_id=thread_id,
+                                                       parent_id=parent_id,
+                                                       conversation_id=conversation_id,
+                                                       content=content,
+                                                       status=status,
+                                                       timestamp=timestamp)
+
+        except Exception as e:
+            logger.error("Error creating system response token message: %s", str(e), exc_info=True)
+            return None
+
+    async def create_system_intermediate_step_message(  # pylint: disable=R0917:too-many-positional-arguments
+        self,
+        message_type: Literal[WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE] = (
+            WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE),
+        message_id: str = str(uuid.uuid4()),
+        thread_id: str = "default",
+        parent_id: str = "default",
+        conversation_id: str | None = None,
+        content: SystemIntermediateStepContent = SystemIntermediateStepContent(name="default", payload="default"),
+        status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
+        timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    ) -> WebSocketSystemIntermediateStepMessage | None:
+        """
+        Creates a system intermediate step message with default values.
+
+        :param message_type: Type of WebSocket message.
+        :param message_id: Unique identifier for the message (default: generated UUID).
+        :param thread_id: ID of the thread the message belongs to (default: "default").
+        :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 status: Status of the message (default: IN_PROGRESS).
+        :param timestamp: Timestamp of the message (default: current UTC time).
+        :return: A WebSocketSystemIntermediateStepMessage instance.
+        """
+        try:
+            return WebSocketSystemIntermediateStepMessage(type=message_type,
+                                                          id=message_id,
+                                                          thread_id=thread_id,
+                                                          parent_id=parent_id,
+                                                          conversation_id=conversation_id,
+                                                          content=content,
+                                                          status=status,
+                                                          timestamp=timestamp)
+
+        except Exception as e:
+            logger.error("Error creating system intermediate step message: %s", str(e), exc_info=True)
+            return None
+
+    async def create_system_interaction_message(  # pylint: disable=R0917:too-many-positional-arguments
+        self,
+        *,
+        message_type: Literal[WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE] = (
+            WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE),
+        message_id: str | None = str(uuid.uuid4()),
+        thread_id: str = "default",
+        parent_id: str = "default",
+        conversation_id: str | None = None,
+        content: HumanPrompt,
+        status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
+        timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    ) -> WebSocketSystemInteractionMessage | None:  # noqa: E125 continuation line with same indent as next logical line
+        """
+        Creates a system interaction message with default values.
+
+        :param message_type: Type of WebSocket message.
+        :param message_id: Unique identifier for the message (default: generated UUID).
+        :param thread_id: ID of the thread the message belongs to (default: "default").
+        :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 status: Status of the message (default: IN_PROGRESS).
+        :param timestamp: Timestamp of the message (default: current UTC time).
+        :return: A WebSocketSystemInteractionMessage instance.
+        """
+        try:
+            return WebSocketSystemInteractionMessage(type=message_type,
+                                                     id=message_id,
+                                                     thread_id=thread_id,
+                                                     parent_id=parent_id,
+                                                     conversation_id=conversation_id,
+                                                     content=content,
+                                                     status=status,
+                                                     timestamp=timestamp)
+
+        except Exception as e:
+            logger.error("Error creating system interaction message: %s", str(e), exc_info=True)
+            return None