jupyterlab
diff --git a/‎.github/workflows/check-release.yml
Lines changed: 1 addition & 1 deletion b/‎.github/workflows/check-release.yml
Lines changed: 1 addition & 1 deletion
diff --git a/‎packages/jupyter-ai-module-cookiecutter/{{cookiecutter.root_dir_name}}/{{cookiecutter.python_name}}/slash_command.py
Lines changed: 2 additions & 2 deletions b/‎packages/jupyter-ai-module-cookiecutter/{{cookiecutter.root_dir_name}}/{{cookiecutter.python_name}}/slash_command.py
Lines changed: 2 additions & 2 deletions
diff --git a/‎packages/jupyter-ai-test/jupyter_ai_test/test_slash_commands.py
Lines changed: 2 additions & 2 deletions b/‎packages/jupyter-ai-test/jupyter_ai_test/test_slash_commands.py
Lines changed: 2 additions & 2 deletions
diff --git a/‎packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py
Lines changed: 5 additions & 3 deletions b/‎packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py
Lines changed: 5 additions & 3 deletions
diff --git a/‎packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py
Lines changed: 16 additions & 13 deletions b/‎packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py
Lines changed: 16 additions & 13 deletions
@@ -24,7 +24,7 @@ jobs:
         uses: jupyter-server/jupyter_releaser/.github/actions/check-release@v2
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
-          version_spec: minor
+          version_spec: "12.34.56"
 
       - name: Upload Distributions
         uses: actions/upload-artifact@v4
 
@@ -1,5 +1,5 @@
 from jupyter_ai.chat_handlers.base import BaseChatHandler, SlashCommandRoutingType
-from jupyter_ai.models import HumanChatMessage
+from jupyterlab_chat.models import Message
 
 
 class TestSlashCommand(BaseChatHandler):
@@ -25,5 +25,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: Message):
         self.reply("This is the `/test` slash command.")
@@ -1,5 +1,5 @@
 from jupyter_ai.chat_handlers.base import BaseChatHandler, SlashCommandRoutingType
-from jupyter_ai.models import HumanChatMessage
+from jupyterlab_chat.models import Message
 
 
 class TestSlashCommand(BaseChatHandler):
@@ -25,5 +25,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: Message):
         self.reply("This is the `/test` slash command.")
@@ -1,9 +1,11 @@
+# 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
 from .default import DefaultChatHandler
-from .export import ExportChatHandler
-from .fix import FixChatHandler
 from .generate import GenerateChatHandler
 from .help import HelpChatHandler
 from .learn import LearnChatHandler
@@ -1,8 +1,8 @@
 import argparse
 from typing import Dict, Type
 
-from jupyter_ai.models import HumanChatMessage
 from jupyter_ai_magics.providers import BaseProvider
+from jupyterlab_chat.models import Message
 from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferWindowMemory
 from langchain_core.prompts import PromptTemplate
@@ -59,7 +59,7 @@ def create_llm_chain(
             verbose=False,
         )
 
-    async def process_message(self, message: HumanChatMessage):
+    async def process_message(self, message: Message):
         args = self.parse_args(message)
         if args is None:
             return
@@ -70,21 +70,24 @@ async def process_message(self, message: HumanChatMessage):
 
         self.get_llm_chain()
 
-        try:
-            with self.pending("Searching learned documents", message):
+        with self.start_reply_stream() as reply_stream:
+            try:
                 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)
-        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)
+
+                # old pending message: "Searching learned documents..."
+                # TODO: configure this pending message in jupyterlab-chat
+                reply_stream.write(response)
+            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.
+                """
+                reply_stream.write(response, message)