fix mypy

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

@@ -315,15 +315,15 @@ def create_llm_chain(
         raise NotImplementedError("Should be implemented by subclasses")
 
     def parse_args(self, message: Message, silent=False):
-        args = message.body.split(" ")
+        args = message.body.split(" ")[1:]
         try:
-            args = self.parser.parse_args(args[1:])
+            arg_namespace = self.parser.parse_args(args[1:])
         except (argparse.ArgumentError, SystemExit) as e:
             if not silent:
                 response = f"{self.parser.format_usage()}"
                 self.reply(response, message)
             return None
-        return args
+        return arg_namespace
 
     def get_llm_chat_memory(self) -> "BaseChatMessageHistory":
         return self.llm_chat_memory

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

@@ -203,7 +203,9 @@ async def process_message(self, message: Message):
         # delete and relearn index if embedding model was changed
         await self.delete_and_relearn()
 
-        with self.pending(f"Loading and splitting files for {load_path}", message):
+        # TODO v3: reinstate pending message
+        # original pending message: "Loading and splitting files for {load_path}"
+        with self.start_reply_stream() as reply_stream:
             try:
                 await self.learn_dir(
                     load_path, args.chunk_size, args.chunk_overlap, args.all_files
@@ -219,7 +221,7 @@ async def process_message(self, message: Message):
                     You can ask questions about these docs by prefixing your message with **/ask**.""" % (
                     load_path.replace("*", r"\*")
                 )
-        self.reply(response, message)
+            reply_stream.write(response)
 
     def _build_list_response(self):
         if not self.metadata.dirs:

packages/jupyter-ai/jupyter_ai/chat_handlers/utils/streaming.py

@@ -71,6 +71,8 @@ def write(self, chunk: str) -> str:
             ),
             append=True,
         )
+
+        return self._stream_id
 
     def close(self):
         self.ychat.awareness.set_local_state_field("isWriting", False)

packages/jupyter-ai/jupyter_ai/history.py

@@ -18,7 +18,7 @@ class YChatHistory(BaseChatMessageHistory):
     TODO: Consider just defining `k` as the number of messages and default to 4.
     """
 
-    def __init__(self, ychat: YChat, k: int = 2):
+    def __init__(self, ychat: YChat, k: Optional[int]):
         self.ychat = ychat
         self.k = k