From 84ebce4ca88fc8b2ed7f4d6e9671edcf743d60c7 Mon Sep 17 00:00:00 2001 From: mc <4072246+mcavdar@users.noreply.github.com> Date: Thu, 31 Oct 2024 19:54:52 +0000 Subject: [PATCH] added AgentStreamMessage to report --- 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):