💬 Stateful Chat

Submitted by Lukas Masuch

Summary

A chat container that automatically keeps track of the chat history.

Functions

`add_message`

Adds a chat message to the chat container.

This command can only be used inside the chat container. The message will be displayed in the UI and added to the chat history so that the same message will be automatically displayed on reruns.

Parameters:

Name	Type	Description	Default
`name`	`Literal['user', 'assistant'] \| str`	The name of the message author. Can be "user" or "assistant" to enable preset styling and avatars. Currently, the name is not shown in the UI but is only set as an accessibility label. For accessibility reasons, you should not use an empty string.	required
`avatar`	`str \| AtomicImage \| None`	The avatar shown next to the message. Can be anything that is supported by the `avatar` parameter of `st.chat_message`. Defaults to None.	`None`
`*args`	`Any`	The content of the message. This can be any number of elements that are supported by `st.write` as well as generator functions to stream content to the UI.	`()`

Raises:

Type	Description
`StreamlitAPIException`	If called outside of a `chat` container.

Source code in src/streamlit_extras/stateful_chat/__init__.py

@extra
def add_message(
    name: str,
    *args: Any,
    avatar: str | AtomicImage | None = None,
) -> None:
    """Adds a chat message to the chat container.

    This command can only be used inside the `chat` container. The message
    will be displayed in the UI and added to the chat history so that the same
    message will be automatically displayed on reruns.

    Args:
        name (Literal["user", "assistant"] | str):
            The name of the message author. Can be "user" or "assistant" to
            enable preset styling and avatars.
            Currently, the name is not shown in the UI but is only set as an
            accessibility label. For accessibility reasons, you should not use
            an empty string.
        avatar (str | AtomicImage | None, optional):
            The avatar shown next to the message. Can be anything that is supported by
            the `avatar` parameter of `st.chat_message`. Defaults to None.
        *args (Any):
            The content of the message. This can be any number of elements that are supported by
            `st.write` as well as generator functions to stream content to the UI.

    Raises:
        StreamlitAPIException: If called outside of a `chat` container.
    """
    active_dg = _active_dg()

    if not hasattr(active_dg, "chat_history"):
        raise StreamlitAPIException("The `add_message` command can only be used inside a `chat` container.")

    # Write to the active container (new_messages_container)
    with active_dg:
        displayed_elements = _display_message(name, *args, avatar=avatar)

    active_dg.chat_history.append(
        ChatMessage(
            author=name,
            avatar=avatar,
            content=displayed_elements,
        )
    )

Import:

from streamlit_extras.stateful_chat import add_message # (1)!

You should add this to the top of your .py file

`chat`

Insert a stateful chat container into your app.

This chat container automatically keeps track of the chat history when you use the add_message command to add messages to the chat.

Parameters:

Name	Type	Description	Default
`key`	`str`	The key that is used to keep track of the chat history in session state. Defaults to "chat_messages".	`'chat_messages'`

Yields:

Name	Type	Description
`DeltaGenerator`	`DeltaGenerator`	The chat container that can be used together with `add_message` to automatically keep track of the chat history.

Source code in src/streamlit_extras/stateful_chat/__init__.py

@extra
@contextmanager
def chat(key: str = "chat_messages") -> Generator[DeltaGenerator, None, None]:
    """Insert a stateful chat container into your app.

    This chat container automatically keeps track of the chat history when you use
    the `add_message` command to add messages to the chat.

    Args:
        key (str, optional): The key that is used to keep track of the chat history in session state.
            Defaults to "chat_messages".

    Yields:
        DeltaGenerator: The chat container that can be used together with `add_message` to
            automatically keep track of the chat history.
    """

    chat_container = st.container()

    if key not in st.session_state:
        st.session_state[key] = []
    chat_history: list[ChatMessage] = st.session_state[key]

    with chat_container:
        # Display existing messages from history
        for message in chat_history:
            _display_message(message["author"], *message["content"], avatar=message["avatar"])

        # Create a container for new messages BEFORE yielding
        # This ensures new messages appear above any content (like chat_input)
        # that the user adds in the yielded block
        new_messages_container = st.container()

    new_messages_container.chat_history = chat_history  # type: ignore

    # Set the active chat container for add_message to use
    token = _active_chat_container.set(new_messages_container)
    try:
        with chat_container:
            yield chat_container
    finally:
        _active_chat_container.reset(token)

Import:

from streamlit_extras.stateful_chat import chat # (1)!

You should add this to the top of your .py file

Examples

`example`

def example() -> None:
    with chat(key="my_chat"):
        if prompt := st.chat_input():
            add_message("user", prompt, avatar="🧑‍💻")

            def stream_echo() -> Generator[str, None, None]:
                for word in prompt.split():
                    yield word + " "
                    time.sleep(0.15)

            add_message("assistant", "Echo: ", stream_echo, avatar="🦜")