diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py index 90bf0aef2..a75f66fff 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py @@ -15,6 +15,7 @@ class AskChatHandler(BaseChatHandler): query the documents from the index, and sends this context to the LLM to generate the final reply. """ + id = "ask-chat-handler" name = "Ask with Local Data" description = "Ask a question augmented with learned data" diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py index 5dde7f6ff..fe0daada5 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py @@ -9,7 +9,6 @@ from jupyter_ai.config_manager import ConfigManager, Logger from jupyter_ai.models import AgentChatMessage, HumanChatMessage from jupyter_ai_magics.providers import BaseProvider - from traitlets.config import Configurable if TYPE_CHECKING: @@ -19,15 +18,17 @@ class BaseChatHandler(Configurable): """Base ChatHandler class containing shared methods and attributes used by multiple chat handler classes.""" - + # Class attributes - id: str = 'base-chat-handler' + id: str = "base-chat-handler" """ID for this chat handler; should be unique""" - name: str = 'Base Chat Handler' # TODO: make NotImplemented + name: str = "Base Chat Handler" # TODO: make NotImplemented """User-facing name of this handler""" - description: str = "Handler for messages that are not commands" # TODO: make NotImplemented + description: str = ( + "Handler for messages that are not commands" # TODO: make NotImplemented + ) """Description used for routing requests, to be used when dispatching messages to model providers. Also shown in the UI.""" diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index 54f42e8c8..e5541dd46 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -1,8 +1,8 @@ -from importlib_metadata import entry_points import logging import time from dask.distributed import Client as DaskClient +from importlib_metadata import entry_points from jupyter_ai.chat_handlers.learn import Retriever from jupyter_ai_magics.utils import get_em_providers, get_lm_providers from jupyter_server.extension.application import ExtensionApp @@ -95,7 +95,7 @@ def initialize_settings(self): "root_dir": self.serverapp.root_dir, "dask_client_future": dask_client_future, # TODO: Set ask_chat_handler based on a retriever related to the learn_chat_handler - "retriever": None + "retriever": None, } for chat_handler_ep in chat_handler_eps: @@ -116,8 +116,10 @@ def initialize_settings(self): command_name = f"/{slash_id}" else: command_name = "default" - self.log.info(f"Trying to register chat handler `{chat_handler.id}` with command `{command_name}`") - + self.log.info( + f"Trying to register chat handler `{chat_handler.id}` with command `{command_name}`" + ) + if command_name in jai_chat_handlers: self.log.error( f"Unable to register chat handler `{chat_handler.id}` because command `{command_name}` already has a handler" @@ -125,8 +127,10 @@ def initialize_settings(self): continue jai_chat_handlers[command_name] = chat_handler - self.log.info(f"Registered chat handler `{chat_handler.id}` with command `{command_name}`.") - + self.log.info( + f"Registered chat handler `{chat_handler.id}` with command `{command_name}`." + ) + self.settings["jai_chat_handlers"] = jai_chat_handlers """