From f8944b1fd32827e2ff125c5b8c5a133dee07f507 Mon Sep 17 00:00:00 2001 From: Mahmut CAVDAR <4072246+mcavdar@users.noreply.github.com> Date: Fri, 1 Nov 2024 22:23:26 +0000 Subject: [PATCH] Fix `/export` by including streamed agent messages (#1077) --- packages/jupyter-ai/jupyter_ai/chat_handlers/export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/export.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/export.py index ed478f57e..7323d81c1 100644 --- 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):