From 3ef11cea7f0795db9df71bd0e190c4e825eccef4 Mon Sep 17 00:00:00 2001
From: apurvakhatri <apurvakhatri2011@gmail.com>
Date: Tue, 27 Feb 2024 12:54:13 -0500
Subject: [PATCH] export chat to markdown functionality added

---
 .../jupyter_ai/chat_handlers/__init__.py      |  1 +
 .../jupyter_ai/chat_handlers/export.py        | 33 +++++++++++++++++++
 packages/jupyter-ai/jupyter_ai/extension.py   |  3 ++
 3 files changed, 37 insertions(+)
 create mode 100644 packages/jupyter-ai/jupyter_ai/chat_handlers/export.py

diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py
index e4c69f012..17364f198 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py
@@ -5,3 +5,4 @@
 from .generate import GenerateChatHandler
 from .help import HelpChatHandler
 from .learn import LearnChatHandler
+from .export import ExportChatHandler
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/export.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/export.py
new file mode 100644
index 000000000..1f5849507
--- /dev/null
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/export.py
@@ -0,0 +1,33 @@
+from typing import List
+
+from jupyter_ai.models import HumanChatMessage, AgentChatMessage
+
+from .base import BaseChatHandler, SlashCommandRoutingType
+
+class ExportChatHandler(BaseChatHandler):
+    id = "export"
+    name = "Export chat messages"
+    help = "Export the chat messages in various formats"
+    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 ""
+
+
+    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
+        with open("./playground/chat_history.md", "w") as chat_history:
+            chat_history.write(markdown_content)
+
+
diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py
index 800c91932..6920792fa 100644
--- a/packages/jupyter-ai/jupyter_ai/extension.py
+++ b/packages/jupyter-ai/jupyter_ai/extension.py
@@ -16,6 +16,7 @@
     GenerateChatHandler,
     HelpChatHandler,
     LearnChatHandler,
+    ExportChatHandler,
 )
 from .chat_handlers.help import HelpMessage
 from .completions.handlers import DefaultInlineCompletionHandler
@@ -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(