Merge pull request #156 from PrefectHQ/streaming

Fix streaming prints
PrefectHQ · Jun 21, 2024 · 496feed · 496feed
2 parents 0171239 + 4cd3b0f
commit 496feed
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/controlflow/llm/formatting.py b/src/controlflow/llm/formatting.py
@@ -9,6 +9,7 @@
 
 from controlflow.llm.messages import (
     AIMessage,
+    AIMessageChunk,
     MessageType,
     ToolMessage,
 )
@@ -36,7 +37,7 @@ def format_message(
     panels = []
     if isinstance(message, ToolMessage):
         return format_tool_message(message, width=width)
-    elif isinstance(message, AIMessage):
+    elif isinstance(message, (AIMessage, AIMessageChunk)):
         if message.content:
             panels.append(format_text_message(message, width=width))