Some typing

jupyterlab · Oct 22, 2024 · 3564bd5 · 3564bd5
1 parent 7ea32a7
commit 3564bd5
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 8 deletions.
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py
@@ -162,7 +162,7 @@ def __init__(
         chat_handlers: Dict[str, "BaseChatHandler"],
         context_providers: Dict[str, "BaseCommandContextProvider"],
         message_interrupted: Dict[str, asyncio.Event],
-        write_message: Callable[[YChat, str], None] | None = None,
+        write_message: Callable[[YChat, str, Optional[str]], None],
     ):
         self.log = log
         self.config_manager = config_manager
@@ -317,7 +317,7 @@ def reply(
         `HumanChatMessage`.
         """
         if chat is not None:
-            self.write_message(chat, response)
+            self.write_message(chat, response, None)
         else:
             agent_msg = AgentChatMessage(
                 id=uuid4().hex,
@@ -507,7 +507,7 @@ def send_help_message(
         )
 
         if chat is not None:
-            self.write_message(chat, help_message_body)
+            self.write_message(chat, help_message_body, None)
         else:
             help_message = AgentChatMessage(
                 id=uuid4().hex,
@@ -518,13 +518,13 @@ def send_help_message(
             )
             self.broadcast_message(help_message)
 
-    def _start_stream(self, human_msg: HumanChatMessage, chat: Optional[YChat]) -> str:
+    def _start_stream(self, human_msg: HumanChatMessage, chat: Optional[YChat]) -> str | None:
         """
         Sends an `agent-stream` message to indicate the start of a response
         stream. Returns the ID of the message, denoted as the `stream_id`.
         """
         if chat is not None:
-            stream_id = self.write_message(chat, "")
+            stream_id = self.write_message(chat, "", None)
         else:
             stream_id = uuid4().hex
             stream_msg = AgentStreamMessage(
@@ -542,7 +542,7 @@ def _start_stream(self, human_msg: HumanChatMessage, chat: Optional[YChat]) -> s
 
     def _send_stream_chunk(
         self,
-        stream_id: str,
+        stream_id: str | None,
         content: str,
         chat: Optional[YChat],
         complete: bool = False,

diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py
@@ -284,7 +284,7 @@ def _add_dir_to_metadata(self, path: str, chunk_size: int, chunk_overlap: int):
             )
         self.metadata.dirs = dirs
 
-    async def delete_and_relearn(self, chat: Optional[YChat]):
+    async def delete_and_relearn(self, chat: Optional[YChat]=None):
         """Delete the vector store and relearn all indexed directories if
         necessary. If the embedding model is unchanged, this method does
         nothing."""

diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py
@@ -2,6 +2,7 @@
 import re
 import time
 import types
+from typing import Optional
 import uuid
 from functools import partial
 
@@ -318,7 +319,7 @@ async def _route(self, message: HumanChatMessage, chat: YChat):
         command_readable = "Default" if command == "default" else command
         self.log.info(f"{command_readable} chat handler resolved in {latency_ms} ms.")
 
-    def write_message(self, chat: YChat, body: str, id: str | None = None) -> str:
+    def write_message(self, chat: YChat, body: str, id: Optional[str]=None) -> str:
         bot = chat.get_user(BOT["username"])
         if not bot:
             chat.set_user(BOT)