[Python] Is there a way to stream ollama chat into a Marimo ui?

I am wondering if there is the possibility of using a stream for Marino notebooks in the ui.chat, something along the lines of:

from ollama import chat
model = 'llama3.2'
def simple_stream(mess, config):
stream = chat(messages=[{'role':'user', 'content':mess}],
model=model,
stream=True)
for chunk in stream:
yield chunk['message']['content']
mo.ui.chat(simple_stream, prompts=[''])


As it is, ui.chat does not really know how to handle the generator. Maybe there is a simple hack that I should consider?

Many thanks for your help!

Continue reading...