From 3564bd54da53114b67a245537ec9a9fa8bac0409 Mon Sep 17 00:00:00 2001
From: Nicolas Brichet <nicolas.brichet@quantstack.net>
Date: Tue, 22 Oct 2024 16:08:09 +0200
Subject: [PATCH] Some typing

---
 packages/jupyter-ai/jupyter_ai/chat_handlers/base.py | 12 ++++++------
 .../jupyter-ai/jupyter_ai/chat_handlers/learn.py     |  2 +-
 packages/jupyter-ai/jupyter_ai/extension.py          |  3 ++-
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py
index 58411ed7..4d2fba09 100644
--- 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
index 7f6ff9b6..a74c147d 100644
--- 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
index 3cf0edb4..9ecdf1cd 100644
--- 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)