Custom write_stream?

[Magnus1]

Is there a way to define a custom st.write_stream()?
I would like a way to stream the response until it reaches a certain keyword or similar logic.

Thanks!

[edsaac]

The original generator could be wrapped in another generator containing the logic that you want to introduce.

Code:

import streamlit as st
from time import sleep
from typing import Generator, Callable


def ai_response():
    """Generator to mimic some LLM response"""
    response = "This is a long AI generated answer, but I wanted it to stop at the first comma :cry:."
    for word in response:
        yield word
        sleep(0.02)


def conditional_writer(func: Callable[[], Generator]):
    """Wrap a generator with some additional logic"""
    generator = func()
    for word in generator:
        if "," in word:
            break
        yield word

    yield "."


def main():
    cols = st.columns(2)

    with cols[0]:
        if st.button("Stream long answer"):
            st.write_stream(ai_response)

    with cols[1]:
        if st.button("Stream short"):
            st.write_stream(conditional_writer(ai_response))


if __name__ == "__main__":
    main()