Fix Tool Call Positioning and Step Loader Animation in Message UI

[kevinwmerritt]

Describe the bug

• Tool calls appear below the result message
• Step loader "skeleton" continues to animate after run is complete
• The result message appears as a second message rather than as part of the initial tool message as the image in the README would suggest.

To Reproduce

All examples use this demo code from the README

import chainlit as cl


@cl.step(type="tool")
async def tool():
    # Fake tool
    await cl.sleep(2)
    return "Response from the tool!"


@cl.on_message  # this function will be called every time a user inputs a message in the UI
async def main(message: cl.Message):
    """
    This function is called every time a user inputs a message in the UI.
    It sends back an intermediate response from the tool, followed by the final answer.

    Args:
        message: The user's message.

    Returns:
        None.
    """

    final_answer = await cl.Message(content="").send()

    # Call the tool
    final_answer.content = await tool()

    await final_answer.update()

Using default config:

cot = "full"

Using default config:

cot = "tool_call"

Expected behavior
Message UI looks like what is described in the documentation here.

Desktop (please complete the following information):

• Mac 14.7 (23H124)
• Chrome 128.0.6613.138 (Official Build)

Additional context
More test cases are needed to to help lock in the expected behavior given the variety of ways to compose steps and messages. The message UI has changed considerably in 2024. I hope it can start to stabilize and have some tests so that regressions can be caught as it adds friction to staying up to date.

Possibly related:

• Incorrect message ordering in UI after 1.1.400 #1230
• Why is the loading icon not showing? #1254
• cl.Step rendering issue #1228

[dosubot]

Hello, @kevinwmerritt! I'm here to help you with any bugs, questions, or contributions. Let's tackle this together!

To fix the issue where tool calls appear below the result message, the step loader continues to animate after the run is complete, and the result message appears as a second message instead of being part of the initial tool message, you need to ensure that the steps are properly managed and updated in the correct order. Here is the corrected code:

1. Ensure that the on_agent_finish method updates the last step correctly.
2. Ensure that the on_tool_finish method updates the tool step correctly and stops the loader.

Here is the updated code:

def on_agent_finish(self, agent_step: AgentStep, **kwargs: Any) -> None:
    if self.last_step:
        run_step = self.last_step
        run_step.end = utc_now()
        run_step.output = agent_step.prompt_node_response
        run_sync(run_step.update())
        self.last_step = None  # Clear the last step after updating

def on_tool_finish(
        self,
        tool_result: str,
        tool_name: Optional[str] = None,
        tool_input: Optional[str] = None,
        **kwargs: Any
    ) -> None:
        # Tool finished, send step with tool_result
        tool_step = self.stack.pop()
        tool_step.output = tool_result
        tool_step.end = utc_now()
        run_sync(tool_step.update())
        self.last_step = tool_step  # Set the last step to the tool step

Additionally, ensure that the frontend component Message.tsx correctly handles the display of messages and loaders. The Message component should be updated to properly manage the state of loaders and message display:

import { keyframes } from '@emotion/react';
import { MessageContext } from 'contexts/MessageContext';
import { memo, useContext } from 'react';

import Box from '@mui/material/Box';
import Stack from '@mui/material/Stack';

import { useConfig } from '@chainlit/react-client';

import { AskUploadButton } from './components/AskUploadButton';
import { MessageAvatar } from './components/Avatar';
import { MessageActions } from './components/MessageActions';
import { MessageButtons } from './components/MessageButtons';
import { MessageContent } from './components/MessageContent';

import { useLayoutMaxWidth } from 'hooks/useLayoutMaxWidth';

import { type IAction, type IMessageElement, type IStep } from 'client-types/';

import { Messages } from './Messages';
import Step from './Step';
import UserMessage from './UserMessage';

interface Props {
  message: IStep;
  elements: IMessageElement[];
  actions: IAction[];
  indent: number;
  isRunning?: boolean;
  isScorable?: boolean;
  scorableRun?: IStep;
}

const Message = memo(
  ({
    message,
    elements,
    actions,
    isRunning,
    indent,
    isScorable,
    scorableRun
  }: Props) => {
    const {
      highlightedMessage,
      defaultCollapseContent,
      allowHtml,
      latex,
      onError
    } = useContext(MessageContext);
    const { config } = useConfig();
    const layoutMaxWidth = useLayoutMaxWidth();
    const isAsk = message.waitForAnswer;
    const isUserMessage = message.type === 'user_message';
    const isStep = !message.type.includes('message');
    // Only keep tool calls if Chain of Thought is tool_call
    const toolCallSkip =
      isStep && config?.ui.cot === 'tool_call' && message.type !== 'tool';

    const hiddenSkip = isStep && config?.ui.cot === 'hidden';

    const skip = toolCallSkip || hiddenSkip;
    if (skip) {
      if (!message.steps) {
        return null;
      }
      return (
        <Messages
          messages={message.steps}
          elements={elements}
          actions={actions}
          indent={indent}
          isRunning={isRunning}
          scorableRun={scorableRun}
        />
      );
    }
    return (
      <>
        <Box
          sx={{
            color: 'text.primary',
            position: 'relative'
          }}
          className="step"
        >
          <Box
            sx={{
              boxSizing: 'border-box',
              mx: 'auto',
              width: '100%',
              maxWidth: indent ? '100%' : layoutMaxWidth,
              px: 2,
              display: 'flex',
              flexDirection: 'column',
              position: 'relative'
            }}
          >
            <Stack
              id={`step-${message.id}`}
              direction="row"
              sx={{
                pb: indent ? 1 : 2,
                flexGrow: 1,
                animation:
                  message.id && highlightedMessage === message.id
                    ? `3s ease-in-out 0.1s ${flash}`
                    : 'none'
              }}
            >
              {/* User message is displayed differently */}
              {isUserMessage ? (
                <Box display="flex" flexDirection="column" flexGrow={1}>
                  <UserMessage message={message}>
                    <MessageContent
                      elements={elements}
                      message={message}
                      preserveSize={
                        !!message.streaming || !defaultCollapseContent
                      }
                      allowHtml={allowHtml}
                      latex={latex}
                    />
                  </UserMessage>
                </Box>
              ) : (
                <Stack
                  direction="row"
                  gap="1rem"
                  width="100%"
                  className="ai-message"
                >
                  {!isStep || !indent ? (
                    <MessageAvatar author={message.name} />
                  ) : null}
                  {/* Display the step and its children */}
                  {isStep ? (
                    <Step step={message} isRunning={isRunning}>
                      {message.steps ? (
                        <Messages
                          messages={message.steps.filter(
                            (s) => !s.type.includes('message')
                          )}
                          elements={elements}
                          actions={actions}
                          indent={indent + 1}
                          isRunning={isRunning}
                        />
                      ) : null}
                      <MessageContent
                        elements={elements}
                        message={message}
                        preserveSize={
                          !!message.streaming || !defaultCollapseContent
                        }
                        allowHtml={allowHtml}
                        latex={latex}
                      />
                      {actions?.length ? (
                        <MessageActions message={message} actions={actions} />
                      ) : null}
                      <MessageButtons message={message} />
                    </Step>
                  ) : (
                    // Display an assistant message
                    <Stack
                      alignItems="flex-start"
                      minWidth={150}
                      flexGrow={1}
                      position="relative"
                    >
                      <MessageContent
                        elements={elements}
                        message={message}
                        preserveSize={
                          !!message.streaming || !defaultCollapseContent
                        }
                        allowHtml={allowHtml}
                        latex={latex}
                      />
                      {!isRunning && isAsk && (
                        <AskUploadButton onError={onError} />
                      )}
                      {actions?.length ? (
                        <MessageActions message={message} actions={actions} />
                      ) : null}
                      <MessageButtons
                        message={message}
                        run={
                          scorableRun && isScorable ? scorableRun : undefined
                        }
                      />
                    </Stack>
                  )}
                </Stack>
              )}
            </Stack>
          </Box>
        </Box>
        {/* Make sure the child assistant messages of a step are displayed at the root level. */}
        {message.steps && isStep ? (
          <Messages
            messages={message.steps.filter((s) => s.type.includes('message'))}
            elements={elements}
            actions={actions}
            indent={0}
            isRunning={isRunning}
            scorableRun={scorableRun}
          />
        ) : null}
        {/* Display the child steps if the message is not a step (usually a user message). */}
        {message.steps && !isStep ? (
          <Messages
            messages={message.steps}
            elements={elements}
            actions={actions}
            indent={indent}
            isRunning={isRunning}
          />
        ) : null}
      </>
    );
  }
);
// Uses yellow[500] with 50% opacity
const flash = keyframes`
  from {
    background-color: transparent;
  }
  25% {
    background-color: rgba(255, 173, 51, 0.5);
  }
  to {
    background-color: transparent;
  }
`;

export { Message };

These changes ensure that the steps are updated in the correct order, the loader stops animating after the run is complete, and the result message appears as part of the initial tool message ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[dokterbob]

@kevinwmerritt Love to hear whether you consider the reply by our beloved bot (which we're still test-driving) relevant or applicable!

Also interested as to the experiences of other users vis-a-vis tool reporting.

@dosu Could you please suggest a more specific/applicable title for the issue? Are there any similar or related issues? If not, rather than relevant to the chainlit library, might it be something about this user's specific use case (do we know enough about their specific use case?)?

[dosubot]

Hello @kevinwmerritt,

To address the issues you're experiencing with the message UI, here are some steps and code adjustments you can make:

1. Tool Call Positioning: Ensure that tool calls are properly integrated into the initial tool message.
2. Step Loader Animation: Make sure the step loader stops animating once the tool run is complete.
3. Message Handling: Ensure the result message updates the initial tool message rather than appearing as a separate message.

Here is the updated code to help fix these issues:

Backend Code Adjustments

Ensure that the on_agent_finish and on_tool_finish methods update the steps correctly:

def on_agent_finish(self, agent_step: AgentStep, **kwargs: Any) -> None:
    if self.last_step:
        run_step = self.last_step
        run_step.end = utc_now()
        run_step.output = agent_step.prompt_node_response
        run_sync(run_step.update())
        self.last_step = None  # Clear the last step after updating

def on_tool_finish(
        self,
        tool_result: str,
        tool_name: Optional[str] = None,
        tool_input: Optional[str] = None,
        **kwargs: Any
    ) -> None:
        # Tool finished, send step with tool_result
        tool_step = self.stack.pop()
        tool_step.output = tool_result
        tool_step.end = utc_now()
        run_sync(tool_step.update())
        self.last_step = tool_step  # Set the last step to the tool step

Frontend Code Adjustments

Ensure that the Message.tsx component correctly handles the display of messages and loaders:

import { keyframes } from '@emotion/react';
import { MessageContext } from 'contexts/MessageContext';
import { memo, useContext } from 'react';

import Box from '@mui/material/Box';
import Stack from '@mui/material/Stack';

import { useConfig } from '@chainlit/react-client';

import { AskUploadButton } from './components/AskUploadButton';
import { MessageAvatar } from './components/Avatar';
import { MessageActions } from './components/MessageActions';
import { MessageButtons } from './components/MessageButtons';
import { MessageContent } from './components/MessageContent';

import { useLayoutMaxWidth } from 'hooks/useLayoutMaxWidth';

import { type IAction, type IMessageElement, type IStep } from 'client-types/';

import { Messages } from './Messages';
import Step from './Step';
import UserMessage from './UserMessage';

interface Props {
  message: IStep;
  elements: IMessageElement[];
  actions: IAction[];
  indent: number;
  isRunning?: boolean;
  isScorable?: boolean;
  scorableRun?: IStep;
}

const Message = memo(
  ({
    message,
    elements,
    actions,
    isRunning,
    indent,
    isScorable,
    scorableRun
  }: Props) => {
    const {
      highlightedMessage,
      defaultCollapseContent,
      allowHtml,
      latex,
      onError
    } = useContext(MessageContext);
    const { config } = useConfig();
    const layoutMaxWidth = useLayoutMaxWidth();
    const isAsk = message.waitForAnswer;
    const isUserMessage = message.type === 'user_message';
    const isStep = !message.type.includes('message');
    // Only keep tool calls if Chain of Thought is tool_call
    const toolCallSkip =
      isStep && config?.ui.cot === 'tool_call' && message.type !== 'tool';

    const hiddenSkip = isStep && config?.ui.cot === 'hidden';

    const skip = toolCallSkip || hiddenSkip;
    if (skip) {
      if (!message.steps) {
        return null;
      }
      return (
        <Messages
          messages={message.steps}
          elements={elements}
          actions={actions}
          indent={indent}
          isRunning={isRunning}
          scorableRun={scorableRun}
        />
      );
    }
    return (
      <>
        <Box
          sx={{
            color: 'text.primary',
            position: 'relative'
          }}
          className="step"
        >
          <Box
            sx={{
              boxSizing: 'border-box',
              mx: 'auto',
              width: '100%',
              maxWidth: indent ? '100%' : layoutMaxWidth,
              px: 2,
              display: 'flex',
              flexDirection: 'column',
              position: 'relative'
            }}
          >
            <Stack
              id={`step-${message.id}`}
              direction="row"
              sx={{
                pb: indent ? 1 : 2,
                flexGrow: 1,
                animation:
                  message.id && highlightedMessage === message.id
                    ? `3s ease-in-out 0.1s ${flash}`
                    : 'none'
              }}
            >
              {/* User message is displayed differently */}
              {isUserMessage ? (
                <Box display="flex" flexDirection="column" flexGrow={1}>
                  <UserMessage message={message}>
                    <MessageContent
                      elements={elements}
                      message={message}
                      preserveSize={
                        !!message.streaming || !defaultCollapseContent
                      }
                      allowHtml={allowHtml}
                      latex={latex}
                    />
                  </UserMessage>
                </Box>
              ) : (
                <Stack
                  direction="row"
                  gap="1rem"
                  width="100%"
                  className="ai-message"
                >
                  {!isStep || !indent ? (
                    <MessageAvatar author={message.name} />
                  ) : null}
                  {/* Display the step and its children */}
                  {isStep ? (
                    <Step step={message} isRunning={isRunning}>
                      {message.steps ? (
                        <Messages
                          messages={message.steps.filter(
                            (s) => !s.type.includes('message')
                          )}
                          elements={elements}
                          actions={actions}
                          indent={indent + 1}
                          isRunning={isRunning}
                        />
                      ) : null}
                      <MessageContent
                        elements={elements}
                        message={message}
                        preserveSize={
                          !!message.streaming || !defaultCollapseContent
                        }
                        allowHtml={allowHtml}
                        latex={latex}
                      />
                      {actions?.length ? (
                        <MessageActions message={message} actions={actions} />
                      ) : null}
                      <MessageButtons message={message} />
                    </Step>
                  ) : (
                    // Display an assistant message
                    <Stack
                      alignItems="flex-start"
                      minWidth={150}
                      flexGrow={1}
                      position="relative"
                    >
                      <MessageContent
                        elements={elements}
                        message={message}
                        preserveSize={
                          !!message.streaming || !defaultCollapseContent
                        }
                        allowHtml={allowHtml}
                        latex={latex}
                      />
                      {!isRunning && isAsk && (
                        <AskUploadButton onError={onError} />
                      )}
                      {actions?.length ? (
                        <MessageActions message={message} actions={actions} />
                      ) : null}
                      <MessageButtons
                        message={message}
                        run={
                          scorableRun && isScorable ? scorableRun : undefined
                        }
                      />
                    </Stack>
                  )}
                </Stack>
              )}
            </Stack>
          </Box>
        </Box>
        {/* Make sure the child assistant messages of a step are displayed at the root level. */}
        {message.steps && isStep ? (
          <Messages
            messages={message.steps.filter((s) => s.type.includes('message'))}
            elements={elements}
            actions={actions}
            indent={0}
            isRunning={isRunning}
            scorableRun={scorableRun}
          />
        ) : null}
        {/* Display the child steps if the message is not a step (usually a user message). */}
        {message.steps && !isStep ? (
          <Messages
            messages={message.steps}
            elements={elements}
            actions={actions}
            indent={indent}
            isRunning={isRunning}
          />
        ) : null}
      </>
    );
  }
);
// Uses yellow[500] with 50% opacity
const flash = keyframes`
  from {
    background-color: transparent;
  }
  25% {
    background-color: rgba(255, 173, 51, 0.5);
  }
  to {
    background-color: transparent;
  }
`;

export { Message };

These changes should help ensure that the steps are updated in the correct order, the loader stops animating after the run is complete, and the result message appears as part of the initial tool message.

For a more specific and applicable title for the issue, consider: "Fix Tool Call Positioning and Step Loader Animation in Message UI".

There are no exact similar threads found, but related issues include:

• "Interaction issues with Chat Profile Description Popover (Resolve #1256 - Fix Interaction Issues with Chat Profile Description Popover #1276)"
• "Centered steps within assistant messages (fix: step should be centered when contained in an assistant message #1324)" ^[1].

These adjustments should help stabilize the message UI and align it with the expected behavior as described in the documentation. If you have any further questions or need additional assistance, feel free to ask!

To continue talking to Dosu, mention @dosu.

[dokterbob]

@kevinwmerritt The centering PR has been recently merged and isn't on main yet. Could try and replicate this issue on main? Perhaps it's already been resolved?

[kevinwmerritt]

Thanks @dokterbob for looking at this.

I believe that @willydouhard had this working properly but it has regressed in recent months. I liked the UI decisions that were made.

Right now when someone new to this project follows the demo in the readme or the documentation they will get a different result than what the screenshot in the readme or docs imply.

This is not a good onboarding experience.

I just would like to see the message UI get back to what @willydouhard intended.

@dokterbob what branch are you referring to?

[kevinwmerritt]

Here is a video of how https://help.chainlit.io/ works with a tool call:

I'm not sure what version help.chainlit.io is currently running, but in 1.2.0 the behavior is different and/or broken.

[dokterbob]

@dokterbob what branch are you referring to?

main

We're aiming for a release on Monday.

[puppetm4st3r]

also happen when you send manual messages from copilot web site to backend and then copilot front render of the message is not in expected order

[willydouhard]

Since the new COT changes, you no longer have to pre send an empty message to display the loader.

This is more natural and will work as intended:

import chainlit as cl


@cl.step(type="tool")
async def tool():
    # Fake tool
    await cl.sleep(2)
    return "Response from the tool!"


@cl.on_message  # this function will be called every time a user inputs a message in the UI
async def main(message: cl.Message):
    """
    This function is called every time a user inputs a message in the UI.
    It sends back an intermediate response from the tool, followed by the final answer.

    Args:
        message: The user's message.

    Returns:
        None.
    """


    # Call the tool
    tool_res = await tool()

    await cl.Message(content=tool_res).send()

I am going to update the docs!

[kevinwmerritt]

Thanks @willydouhard. That makes a lot of sense.

Here is what I'm seeing now:

• When using the step decorator, the tool message does not "complete" in the UI.
• It properly says "Using " but does not say "Used " at the end.

This is with 1.2.0