diff --git a/src/funcchain/streaming.py b/src/funcchain/streaming.py index 163a9c8..34b3c68 100644 --- a/src/funcchain/streaming.py +++ b/src/funcchain/streaming.py @@ -85,9 +85,12 @@ def stream_to( >>> with stream_to(print): ... # your chain calls here """ + import builtins, rich + # TODO: set streaming to true specifically in this context - if fn is print and kwargs == {}: + if (fn is builtins.print or fn is rich.print) and kwargs == {}: kwargs = {"end": "", "flush": True} + cb = AsyncStreamHandler(fn, kwargs) stream_handler.set(cb) yield cb