feat(Chat): Add .message_stream_context()

[cpsievert]

This PR adds a .message_stream_context() method to Chat, a new more composable and configurable way to append streaming messages to the chat. This context manager can be used in isolation, and also nested within itself, which is primarily useful making checkpoints in the message stream to .clear() back to:

Basic example app.py

import asyncio

from shiny import reactive
from shiny.express import ui

chat = ui.Chat(id="my_chat")
chat.ui()

@reactive.effect
async def _():
    async with chat.message_stream_context() as outer:
        await outer.append("Starting stream...\n\nProgress:")
        async with chat.message_stream_context() as inner:
            for x in [0, 50, 100]:
                await inner.append(f" {x}%")
                await asyncio.sleep(1)
                await inner.clear()
        await outer.clear()
        await outer.append("Completed stream")

Kapture.2025-03-17.at.11.18.33.mp4

A big motivator for adding .message_stream_context() is to have a mechanism to display content during tool calls. In this context, you likely have a long-running .append_message_stream() handling the response generation, but need to also display other UI in the message during a tool call. Here's an mock example of what that might look like:

Tool call app.py

import asyncio

from shiny import reactive
from shiny.express import input, render, ui

ui.page_opts(title="Hello message streams")

chat = ui.Chat(id="chat")
chat.ui()

# Launch a stream on load
@reactive.effect
async def _():
    await chat.append_message_stream(mock_stream())

async def mock_stream():
    yield "Starting outer stream...\n\n"
    await asyncio.sleep(0.5)
    await mock_tool()
    await asyncio.sleep(0.5)
    yield "\n\n...outer stream complete"

async def mock_tool():
    steps = [
        "Starting inner stream 🔄...\n\n",
        "Progress: 0%...",
        "Progress: 50%...",
        "Progress: 100%...",
    ]
    async with chat.message_stream_context() as msg:
        for chunk in steps:
            await msg.append(chunk)
            await asyncio.sleep(0.5)
        await msg.clear()
        await msg.append("Completed inner stream ✅")

Kapture.2025-03-17.at.11.28.19.mp4

This PR also comes with a performance bonus that streaming messages no longer send the entire accumulated content string on every update (that will still be necessary, however, when @transform_assistant_response is used). I figured it was addressing in this PR, since it took some surgery to get ._append_message_chunk() correctly reflected in the server-side message state to reflected.

[jcheng5]

I don't love operation = "replace" as part of the public API (even though I added it--I was intending for it to be a detail of the implementation/protocol). And I also am not in love with .message_stream() as a name, both of those terms are present in so many different methods at this point that it's hard to know what it means.

What if we solve both problems, by replacing:

with chat.message_stream():
    chat.append_message_chunk("Foo")
    chat.append_message_chunk("Bar")
    chat.append_message_chunk("Baz", operation="replace")

with:

with chat.message_checkpoint() as cp:
    chat.append_message_chunk("Foo")
    chat.append_message_chunk("Bar")
    cp.restore()
    chat.append_message_chunk("Baz")

We wouldn't even need the _message_stream_checkpoint member anymore, as calling message_checkpoint() would return an object that contains a snapshot of _current_stream_message; calling cp.restore() would simply be the equivalent of .append_message_chunk(operation = "replace"). So in fact, it doesn't even need to be a context manager:

cp = chat.message_checkpoint()
chat.append_message_chunk("Foo")
chat.append_message_chunk("Bar")
cp.restore()
chat.append_message_chunk("Baz")

In theory maybe it's slightly less efficient because the restoring and the sending of "Baz" happens as two separate messages instead of one, but we're talking about a scenario where each token is generally being streamed as a separate message already, so who cares.

The other downside (if it's even a downside, arguably it's a good thing) is that in order to restore to any checkpoint, you need to actually have the instance of that checkpoint. As opposed to the current approach, where you can only restore to the innermost checkpoint, but you can do it from anywhere without anything but the chat object itself.

[cpsievert]

One thing to consider is that the current implementation of chat.message_stream() allows you to create a message stream in isolation, which is pretty nice for computing on a stream:

@chat.on_user_submit
async def _(user_input):
    stream = await chat_model.stream_async(user_input)
    with chat.message_stream():
        for chunk in stream:
            chat.append_message_chunk(chunk.upper())

So, given that we like that behavior, and it's worth keeping, I'm not in love with chat.message_checkpoint() as a name. I do agree .message_stream() is confusingly close to .append_message_stream(), but I'm also not sure that's purely a bad thing since it is an alternative to .append_message_stream().

That said, I like the .restore() idea, so I'll be for sure adding that, but will have to stew on the name/API a bit more.

[cpsievert]

@jcheng5 I'm now leaning towards .append_message_context() for a name (see the updated #1906 (comment)). I also got rid of .append_message_chunk() to help reduce name clutter. Let me know what you think (even if it's: "I still standby my original suggestion").

Also tagging in @gadenbuie to see what he thinks.

[gadenbuie]

I really liked with chat.message_stream() as a name. Reading through the above, my feeling is that something like chat.message_checkpoint() is probably accurate but follows the implementation too closely (i.e. we could come up with a different approach that does the same thing and doesn't rely on checkpointing).

I do not like with chat.append_message_context() for a few reasons, but mostly because I think context managers should following the pattern: with noun do verbs. But append_message_context() sounds like it would append context to a message, not create a new message context for appending?

I think with chat.message_stream_context() would be clearer if we're trying to find something more specific than .message_stream().

[jcheng5]

If you're not going to use the term "checkpoint" then I think this should be msg.clear(). What do you think?

[cpsievert]

Yea, that is better, thanks 👍

[cpsievert]

After playing with this with some real world examples, I feel pretty strongly that we shouldn't have a .clear() at all, just a .replace(). That's because, if you don't do the update in one-shot, it's too easy to get a flickering effect.

Will be push an update soon