Fix /export by including streamed agent messages (#1077)

jupyterlab · Nov 1, 2024 · f8944b1 · f8944b1
1 parent da1f5f6
commit f8944b1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/export.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/export.py
@@ -3,7 +3,7 @@
 from datetime import datetime
 from typing import List
 
-from jupyter_ai.models import AgentChatMessage, HumanChatMessage
+from jupyter_ai.models import AgentChatMessage, AgentStreamMessage, HumanChatMessage
 
 from .base import BaseChatHandler, SlashCommandRoutingType
 
@@ -22,7 +22,7 @@ def __init__(self, *args, **kwargs):
         self.parser.add_argument("path", nargs=argparse.REMAINDER)
 
     def chat_message_to_markdown(self, message):
-        if isinstance(message, AgentChatMessage):
+        if isinstance(message, (AgentChatMessage, AgentStreamMessage)):
             agent = self.config_manager.persona.name
             return f"**{agent}**: {message.body}"
         elif isinstance(message, HumanChatMessage):