on_chat_resume function is unable to retrieve the metadata of messages when using the thread's steps

[constantinidan]

Describe the bug
The on_chat_resume function is unable to retrieve the metadata of messages when using the thread's steps.

To Reproduce
Start a chat session and send messages.
Resume the chat session.
In the on_chat_resume function, try to access message metadata using thread["steps"].

Expected behavior
When resuming a chat, the on_chat_resume function should be able to retrieve the full message history, including metadata, from the thread["steps"] object. Each message in the steps should contain its associated metadata.
Currently, the code uses a workaround by storing the chat history in the thread metadata (thread["metadata"]["chat_history"]). While this works, it would be more efficient and consistent to retrieve the message history directly from the thread["steps"] object, including all associated metadata.

[constantinidan]

Steps to reproduce

from operator import itemgetter
import os
from mistralai import Mistral
from dotenv import load_dotenv

from chainlit.types import ThreadDict
import chainlit as cl

load_dotenv()

@cl.password_auth_callback
def auth():
    return cl.User(identifier="test")


@cl.on_chat_start
async def on_chat_start():
    cl.user_session.set("chat_history", [])


@cl.on_chat_resume
async def on_chat_resume(thread: ThreadDict):
    cl.user_session.set("chat_history", [])

    # Option 1: use the chat history saved in the thread metadata. Which will contain the metadata of each message
    chat_history = thread["metadata"]["chat_history"]
    for message in chat_history:
        print(message)
        cl.user_session.get("chat_history").append(message)

    # # Option 2: use the steps to fetch the chat history. But there is a bug: existing metadata is not returned
    # for message in thread["steps"]:
    #     # BUG: the steps do not return metadata
    #     if message["type"] == "user_message":
    #         cl.user_session.get("chat_history").append({"role": "user", "content": message["output"]})
    #     elif message["type"] == "assistant_message":
    #         cl.user_session.get("chat_history").append({"role": "assistant", "content": message["output"]})


@cl.on_message
async def on_message(message: cl.Message):
    # Note: by default, the list of messages is saved and the entire user session is saved in the thread metadata
    chat_history = cl.user_session.get("chat_history")

    api_key = os.environ["MISTRAL_API_KEY"]
    model = "ministral-8b-latest"

    client = Mistral(api_key=api_key)

    chat_history.append({"role": "user", "content": message.content})

    chat_response = client.chat.complete(
        model=model,
        messages=chat_history
    )
    
    response_content = chat_response.choices[0].message.content

    # workaround to save metadata and fetch it in the resume on_chat_resume method
    chat_history.append({"role": "assistant", "content": response_content, "metadata": {"test": "test"}})

    # The metadata will be saved in the database. But not fetched in the resume on_chat_resume method, cf bug explained above
    await cl.Message(content=response_content, metadata={"test": "test"}).send()

[dosubot]

The issue you're encountering is that the on_chat_resume function is unable to retrieve message metadata from the thread["steps"] object. This is because the metadata is not being included in the steps when the chat is resumed.

In Chainlit, the steps attribute in the thread object is populated by converting each step into a StepDict, which should include metadata if it is present in the original step object ^[1][2]. However, it seems that in your case, the metadata is not being correctly included or retrieved from the steps.

As a workaround, you are currently storing the chat history, including metadata, in the thread["metadata"]["chat_history"], which allows you to access the metadata correctly. This approach works but is not as efficient as directly accessing the metadata from thread["steps"].

To address this issue, you might need to ensure that the metadata is correctly set and retrieved in the StepDict conversion process. This could involve checking the implementation of how steps are converted and stored, ensuring that metadata is not being inadvertently omitted during this process. If the issue persists, it might require further investigation into how the data layer handles the persistence and retrieval of step metadata.

To continue talking to Dosu, mention @dosu.

Written by @dosu, approved by dokterbob