Use jupyterlab-chat extension

packages/jupyter-ai-module-cookiecutter/{{cookiecutter.root_dir_name}}/{{cookiecutter.python_name}}/slash_command.py

@@ -1,5 +1,6 @@
 from jupyter_ai.chat_handlers.base import BaseChatHandler, SlashCommandRoutingType
 from jupyter_ai.models import HumanChatMessage
+from jupyterlab_chat.ychat import YChat
 
 
 class TestSlashCommand(BaseChatHandler):
@@ -25,5 +26,5 @@ class TestSlashCommand(BaseChatHandler):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-    async def process_message(self, message: HumanChatMessage):
+    async def process_message(self, message: HumanChatMessage, chat: YChat):
         self.reply("This is the `/test` slash command.")

packages/jupyter-ai-test/jupyter_ai_test/test_slash_commands.py

@@ -1,5 +1,6 @@
 from jupyter_ai.chat_handlers.base import BaseChatHandler, SlashCommandRoutingType
 from jupyter_ai.models import HumanChatMessage
+from jupyterlab_chat.ychat import YChat
 
 
 class TestSlashCommand(BaseChatHandler):
@@ -25,5 +26,5 @@ class TestSlashCommand(BaseChatHandler):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-    async def process_message(self, message: HumanChatMessage):
-        self.reply("This is the `/test` slash command.")
+    async def process_message(self, message: HumanChatMessage, chat: YChat):
+        self.reply("This is the `/test` slash command.", chat)

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

@@ -1,3 +1,8 @@
+# The following import is to make sure jupyter_ydoc is imported before
+# jupyterlab_chat, otherwise it leads to circular import because of the
+# YChat relying on YBaseDoc, and jupyter_ydoc registering YChat from the entry point.
+import jupyter_ydoc
+
 from .ask import AskChatHandler
 from .base import BaseChatHandler, SlashCommandRoutingType
 from .clear import ClearChatHandler

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

@@ -1,8 +1,9 @@
 import argparse
-from typing import Dict, Type
+from typing import Dict, Optional, Type
 
 from jupyter_ai.models import HumanChatMessage
 from jupyter_ai_magics.providers import BaseProvider
+from jupyterlab_chat.ychat import YChat
 from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferWindowMemory
 from langchain_core.prompts import PromptTemplate
@@ -59,32 +60,32 @@ def create_llm_chain(
             verbose=False,
         )
 
-    async def process_message(self, message: HumanChatMessage):
-        args = self.parse_args(message)
+    async def process_message(self, message: HumanChatMessage, chat: Optional[YChat]):
+        args = self.parse_args(message, chat)
         if args is None:
             return
         query = " ".join(args.query)
         if not query:
-            self.reply(f"{self.parser.format_usage()}", message)
+            self.reply(f"{self.parser.format_usage()}", chat, message)
             return
 
         self.get_llm_chain()
 
         try:
-            with self.pending("Searching learned documents", message):
+            with self.pending("Searching learned documents", message, chat=chat):
                 assert self.llm_chain
                 # TODO: migrate this class to use a LCEL `Runnable` instead of
                 # `Chain`, then remove the below ignore comment.
                 result = await self.llm_chain.acall(  # type:ignore[attr-defined]
                     {"question": query}
                 )
                 response = result["answer"]
-            self.reply(response, message)
+            self.reply(response, chat, message)
         except AssertionError as e:
             self.log.error(e)
             response = """Sorry, an error occurred while reading the from the learned documents.
             If you have changed the embedding provider, try deleting the existing index by running
             `/learn -d` command and then re-submitting the `learn <directory>` to learn the documents,
             and then asking the question again.
             """
-            self.reply(response, message)
+            self.reply(response, chat, message)