From 5d89204e3700af1c3d6dfbed30ff5e6bc53a7a98 Mon Sep 17 00:00:00 2001 From: "David L. Qiu" Date: Wed, 4 Dec 2024 11:00:29 -0800 Subject: [PATCH] pre-commit --- .../jupyter_ai/chat_handlers/base.py | 20 ++++++------------- .../jupyter_ai/chat_handlers/fix.py | 4 +--- .../jupyter_ai/chat_handlers/generate.py | 4 +--- .../jupyter_ai/chat_handlers/learn.py | 8 ++------ packages/jupyter-ai/jupyter_ai/extension.py | 18 +++++++++-------- .../jupyter_ai/tests/test_handlers.py | 2 +- 6 files changed, 21 insertions(+), 35 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py index c2e1a86d7..233099151 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py @@ -256,9 +256,7 @@ async def process_message(self, message: HumanChatMessage): """ raise NotImplementedError("Should be implemented by subclasses.") - async def handle_exc( - self, e: Exception, message: HumanChatMessage - ): + async def handle_exc(self, e: Exception, message: HumanChatMessage): """ Handles an exception raised by `self.process_message()`. A default implementation is provided, however chat handlers (subclasses) should @@ -266,9 +264,7 @@ async def handle_exc( """ await self._default_handle_exc(e, message) - async def _default_handle_exc( - self, e: Exception, message: HumanChatMessage - ): + async def _default_handle_exc(self, e: Exception, message: HumanChatMessage): """ The default definition of `handle_exc()`. This is the default used when the `handle_exc()` excepts. @@ -285,13 +281,13 @@ async def _default_handle_exc( f"Sorry, an error occurred. Details below:\n\n```\n{formatted_e}\n```" ) self.reply(response, message) - + def write_message(self, body: str, id: Optional[str] = None) -> None: """[Jupyter Chat only] Writes a message to the YChat shared document that this chat handler is assigned to.""" if not self.ychat: return - + bot = self.ychat.get_user(BOT["username"]) if not bot: self.ychat.set_user(BOT) @@ -422,9 +418,7 @@ def pending( TODO: Simplify it by only modifying the awareness as soon as jupyterlab chat is the only used chat. """ - pending_msg = self.start_pending( - text, human_msg=human_msg, ellipsis=ellipsis - ) + pending_msg = self.start_pending(text, human_msg=human_msg, ellipsis=ellipsis) try: yield pending_msg finally: @@ -501,9 +495,7 @@ def output_dir(self) -> str: else: return self.root_dir - def send_help_message( - self, human_msg: Optional[HumanChatMessage] = None - ) -> None: + def send_help_message(self, human_msg: Optional[HumanChatMessage] = None) -> None: """Sends a help message to all connected clients.""" lm_provider = self.config_manager.lm_provider unsupported_slash_commands = ( diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py index bb271ea3a..390b93cf6 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py @@ -103,6 +103,4 @@ async def process_message(self, message: HumanChatMessage): "error_name": selection.error.name, "error_value": selection.error.value, } - await self.stream_reply( - inputs, message, pending_msg="Analyzing error" - ) + await self.stream_reply(inputs, message, pending_msg="Analyzing error") diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py index 1ae4d0bf7..a69b5ed28 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py @@ -273,9 +273,7 @@ async def process_message(self, message: HumanChatMessage): response = f"""🎉 I have created your notebook and saved it to the location {final_path}. I am still learning how to create notebooks, so please review all code before running it.""" self.reply(response, message) - async def handle_exc( - self, e: Exception, message: HumanChatMessage - ): + async def handle_exc(self, e: Exception, message: HumanChatMessage): timestamp = time.strftime("%Y-%m-%d-%H.%M.%S") default_log_dir = Path(self.output_dir) / "jupyter-ai-logs" log_dir = self.log_dir or default_log_dir diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py index 8e0062add..c350dd1b8 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py @@ -143,9 +143,7 @@ async def process_message(self, message: HumanChatMessage): if args.delete: self.delete() - self.reply( - f"👍 I have deleted everything I previously learned.", message - ) + self.reply(f"👍 I have deleted everything I previously learned.", message) return if args.list: @@ -205,9 +203,7 @@ async def process_message(self, message: HumanChatMessage): # 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 - ): + with self.pending(f"Loading and splitting files for {load_path}", message): try: await self.learn_dir( load_path, args.chunk_size, args.chunk_overlap, args.all_files diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index 6f0499016..5080fc4bd 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -6,6 +6,7 @@ from functools import partial from typing import Dict, Optional +import traitlets from dask.distributed import Client as DaskClient from importlib_metadata import entry_points from jupyter_ai.chat_handlers.learn import Retriever @@ -18,7 +19,6 @@ from jupyterlab_chat.ychat import YChat from pycrdt import ArrayEvent from tornado.web import StaticFileHandler -import traitlets from traitlets import Integer, List, Unicode from .chat_handlers import ( @@ -237,7 +237,7 @@ def initialize(self): """ Nested dictionary that returns the dedicated chat handler instance that should be used, given the room ID and command ID respectively. - + Example: `self.chat_handlers_by_room[]` yields the set of chat handlers dedicated to the room identified by ``. """ @@ -249,11 +249,11 @@ def initialize(self): self.event_logger.add_listener( schema_id=JUPYTER_COLLABORATION_EVENTS_URI, listener=self.connect_chat ) - + async def connect_chat( self, logger: EventLogger, schema_id: str, data: dict ) -> None: - # ignore events that are not chat room initialization events + # ignore events that are not chat room initialization events if not ( data["room"].startswith("text:chat:") and data["action"] == "initialize" @@ -276,7 +276,7 @@ async def connect_chat( ) if ychat.awareness is not None: ychat.awareness.set_local_state_field("user", BOT) - + # initialize chat handlers for new chat self.chat_handlers_by_room[room_id] = self._init_chat_handlers(ychat) @@ -504,11 +504,13 @@ async def _stop_extension(self): await dask_client.close() self.log.debug("Closed Dask client.") - def _init_chat_handlers(self, ychat: Optional[YChat] = None) -> Dict[str, BaseChatHandler]: + def _init_chat_handlers( + self, ychat: Optional[YChat] = None + ) -> Dict[str, BaseChatHandler]: """ Initializes a set of chat handlers. May accept a YChat instance for collaborative chats. - + TODO: Make `ychat` required once Jupyter Chat migration is complete. """ eps = entry_points() @@ -528,7 +530,7 @@ def _init_chat_handlers(self, ychat: Optional[YChat] = None) -> Dict[str, BaseCh "chat_handlers": chat_handlers, "context_providers": self.settings["jai_context_providers"], "message_interrupted": self.settings["jai_message_interrupted"], - "ychat": ychat + "ychat": ychat, } default_chat_handler = DefaultChatHandler(**chat_handler_kwargs) clear_chat_handler = ClearChatHandler(**chat_handler_kwargs) diff --git a/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py b/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py index d869c3118..4d851b2b2 100644 --- a/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py +++ b/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py @@ -78,7 +78,7 @@ def broadcast_message(message: Message) -> None: chat_handlers={}, context_providers={}, message_interrupted={}, - ychat=None + ychat=None, )