/export added

packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py

@@ -2,6 +2,7 @@
 from .base import BaseChatHandler, SlashCommandRoutingType
 from .clear import ClearChatHandler
 from .default import DefaultChatHandler
+from .export import ExportChatHandler
 from .generate import GenerateChatHandler
 from .help import HelpChatHandler
 from .learn import LearnChatHandler

packages/jupyter-ai/jupyter_ai/chat_handlers/export.py

@@ -0,0 +1,46 @@
+import os
+from typing import List
+
+from jupyter_ai.models import AgentChatMessage, HumanChatMessage
+
+from .base import BaseChatHandler, SlashCommandRoutingType
+
+
+class ExportChatHandler(BaseChatHandler):
+    id = "export"
+    name = "Export chat messages"
+    help = "Export the chat messages in markdown format"
+    routing_type = SlashCommandRoutingType(slash_id="export")
+
+    uses_llm = False
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def chat_message_to_markdown(self, message):
+        if isinstance(message, AgentChatMessage):
+            return f"**Agent**: {message.body}"
+        elif isinstance(message, HumanChatMessage):
+            return f"**{message.client.display_name}**: {message.body}"
+        else:
+            return ""
+
+    # Multiple chat histories in separate files
+    def get_chat_filename(self, path="./chat_history.md"):
+        filename, extension = os.path.splitext(path)
+        counter = 1
+        while os.path.exists(path):
+            path = filename + "_" + str(counter) + ".md"
+            counter += 1
+        return path
+
+    async def process_message(self, _):
+        markdown_content = "\n\n".join(
+            self.chat_message_to_markdown(msg) for msg in self._chat_history
+        )
+        # Write the markdown content to a file or do whatever you want with it
+        chat_filename = self.get_chat_filename()
+        with open(chat_filename, "w") as chat_history:
+            chat_history.write(markdown_content)
+
+        self.reply(f"File saved to `{chat_filename}`")

packages/jupyter-ai/jupyter_ai/extension.py

@@ -13,6 +13,7 @@
     AskChatHandler,
     ClearChatHandler,
     DefaultChatHandler,
+    ExportChatHandler,
     GenerateChatHandler,
     HelpChatHandler,
     LearnChatHandler,
@@ -239,12 +240,14 @@ def initialize_settings(self):
         retriever = Retriever(learn_chat_handler=learn_chat_handler)
         ask_chat_handler = AskChatHandler(**chat_handler_kwargs, retriever=retriever)
 
+        export_chat_handler = ExportChatHandler(**chat_handler_kwargs)
         jai_chat_handlers = {
             "default": default_chat_handler,
             "/ask": ask_chat_handler,
             "/clear": clear_chat_handler,
             "/generate": generate_chat_handler,
             "/learn": learn_chat_handler,
+            "/export": export_chat_handler,
         }
 
         help_chat_handler = HelpChatHandler(