Chainlit: Trying to nest AskUserMessage under a step

chainlit-nested-ask.py

from typing import Any, Callable, cast, Optional, Union
import textwrap
import chainlit as cl
from datetime import datetime
from chainlit.sync import run_sync
from chainlit.message import AskMessageBase, AskUserMessage, MessageStepType
from chainlit.types import AskSpec
from chainlit.step import StepDict
from chainlit.config import config
from chainlit.telemetry import trace_event
from chainlit.context import context

class AskUserMessageNested(AskMessageBase):
    """
    Patch of AskUserMessage that takes a `parent_id` in constructor, and
    sticks it into the `step_dict` to force the prompt to be nested under a step.
    The intent is that both the prompt and user answer should be nested under a step,
    but this only works partially -- i.e. the prompt does get nested, but the user answer
    is immediately echoed at root level. One of two things might help with nesting the user's 
    answer under the step:
    - a way to have the echoing of user input be under the given parent_id
    - a way to delete an already-echoed root-level user input, so that we can capture
      the user's response and render it ourselves under the step.
    """

    def __init__(
        self,
        content: str,
        author: str = config.ui.name,
        type: MessageStepType = "assistant_message",
        disable_feedback: bool = False,
        timeout: int = 60,
        raise_on_timeout: bool = False,
        parent_id: Optional[str] = None,
    ):
        self.content = content
        self.author = author
        self.timeout = timeout
        self.type = type
        self.disable_feedback = disable_feedback
        self.raise_on_timeout = raise_on_timeout
        self.parent_id = parent_id

        super().__post_init__()

    async def send(self) -> Union[StepDict, None]:
        """
        Sends the question to ask to the UI and waits for the reply.
        """
        trace_event("send_ask_user")
        if not self.created_at:
            self.created_at = datetime.utcnow().isoformat()

        if config.code.author_rename:
            self.author = await config.code.author_rename(self.author)

        if self.streaming:
            self.streaming = False

        self.wait_for_answer = True

        step_dict = await self._create()
        if self.parent_id:
            step_dict["parentId"] = self.parent_id

        spec = AskSpec(type="text", timeout=self.timeout)

        res = cast(
            Union[None, StepDict],
            await context.emitter.send_ask_user(step_dict, spec, self.raise_on_timeout),
        )

        self.wait_for_answer = False

        return res
      

import chainlit as cl


@cl.on_chat_start
async def main():
    step = cl.Step(name="Root", type="step")
    step.output = "Say something"
    await step.send()

    # ask user's name, NESTED under this step.
    # The QUESTION does get nested, but user's answer is immediately 
    # echoed at Root level, not under the step.
    
    res = await AskUserMessageNested(
        content="What is your name?",
        timeout=60_000,
        parent_id=step.id, # so it is nested under the step
    ).send()