dosubot[bot]
commented
1 day ago 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:
- Ensure that the
on_agent_finishmethod updates the last step correctly. - Ensure that the
on_tool_finishmethod 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 stepAdditionally, 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
Describe the bug
To Reproduce
All examples use this demo code from the README
```python 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:
Using default config:
Expected behavior Message UI looks like what is described in the documentation here.
Desktop (please complete the following information):
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: