Skip to content

Commit

Permalink
/export added (#658)
Browse files Browse the repository at this point in the history
* export chat to markdown functionality added

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update export.py

- Added multiple chat history files
- changed export to mention markdown files only
- respond with filename in chat

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Edits to /export function

Using Markdown formatting when displaying file paths

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Sanjiv Das <[email protected]>
  • Loading branch information
3 people authored Mar 18, 2024
1 parent b18b7d0 commit 4814904
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 0 deletions.
1 change: 1 addition & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
46 changes: 46 additions & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/export.py
Original file line number Diff line number Diff line change
@@ -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}`")
3 changes: 3 additions & 0 deletions packages/jupyter-ai/jupyter_ai/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
AskChatHandler,
ClearChatHandler,
DefaultChatHandler,
ExportChatHandler,
GenerateChatHandler,
HelpChatHandler,
LearnChatHandler,
Expand Down Expand Up @@ -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(
Expand Down

0 comments on commit 4814904

Please sign in to comment.