WIP: Updates per feedback, adds custom handler

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

@@ -1,6 +1,7 @@
 from .ask import AskChatHandler
 from .base import BaseChatHandler
 from .clear import ClearChatHandler
+from .custom import CustomChatHandler
 from .default import DefaultChatHandler
 from .generate import GenerateChatHandler
 from .help import HelpChatHandler

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

@@ -26,10 +26,10 @@ class AskChatHandler(BaseChatHandler):
     to the LLM to generate the final reply.
     """
 
-    id = "ask-chat-handler"
+    id = "ask"
     name = "Ask with Local Data"
-    description = "Ask a question augmented with learned data"
     help = "Asks a question with retrieval augmented generation (RAG)"
+    routing_method = "slash_command"
     slash_id = "ask"
 
     def __init__(self, retriever, *args, **kwargs):

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

@@ -3,7 +3,7 @@
 import traceback
 
 # necessary to prevent circular import
-from typing import TYPE_CHECKING, Dict, Optional, Type
+from typing import TYPE_CHECKING, ClassVar, Dict, Optional, Type
 from uuid import uuid4
 
 from jupyter_ai.config_manager import ConfigManager, Logger
@@ -20,27 +20,30 @@ class BaseChatHandler(Configurable):
     multiple chat handler classes."""
 
     # Class attributes
-    id: str = "base-chat-handler"
+    id: ClassVar[str] = "base"
     """ID for this chat handler; should be unique"""
 
-    name: str = "Base Chat Handler"  # TODO: make NotImplemented
+    name: ClassVar[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: ClassVar[Optional[str]]
     """Description used for routing requests, to be used when dispatching
-    messages to model providers. Also shown in the UI."""
+    messages to model providers. Also shown in the UI for transparency;
+    optimized for model interpretation, not human-facing help.
+    Not necessary when the routing method is "slash_command"."""
 
     # TODO: make NotImplemented
-    help: str = "This is used when the message in the chat interface is not a command"
+    help: ClassVar[str] = "This is used when the message in the chat interface is not a command"
     """What this chat handler does, which third-party models it contacts,
     the format of the data it returns to the user, etc. Used in the UI."""
 
-    slash_id: Optional[str]
+    # TODO: make NotImplemented
+    routing_method: ClassVar[str] = "slash_command"
+
+    slash_id: ClassVar[Optional[str]]
     """Slash ID for routing a chat command to this handler. Only one handler
     may declare a particular slash ID. Must contain only alphanumerics and
-    underscores."""
+    underscores. Must not be specified if routing method is "natural_language"."""
 
     def __init__(
         self,

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

@@ -6,10 +6,10 @@
 
 
 class ClearChatHandler(BaseChatHandler):
-    id = "clear-chat-handler"
+    id = "clear"
     name = "Clear chat messages"
-    description = "Clear the chat message history display"
     help = "Clears the displayed chat message history only; does not clear the context sent to chat providers"
+    routing_method = "slash_command"
     slash_id = "clear"
 
     def __init__(self, chat_history: List[ChatMessage], *args, **kwargs):

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

@@ -0,0 +1,33 @@
+import time
+from typing import List
+from uuid import uuid4
+
+from jupyter_ai.models import AgentChatMessage, HumanChatMessage
+
+from .base import BaseChatHandler
+
+CUSTOM_MESSAGE = "This handler displays a custom message in response to any prompt."
+
+def CustomMessage():
+    return AgentChatMessage(
+        id=uuid4().hex,
+        time=time.time(),
+        body=CUSTOM_MESSAGE,
+        reply_to="",
+    )
+
+"""
+This is a sample custom chat handler class to demonstrate entry points.
+"""
+class CustomChatHandler(BaseChatHandler):
+    id = "custom"
+    name = "Custom"
+    help = "Displays a custom message in the chat message area"
+    routing_method = "slash_command"
+    slash_id = "custom"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    async def _process_message(self, message: HumanChatMessage):
+        self.reply(CUSTOM_MESSAGE, message)

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

@@ -213,10 +213,10 @@ def create_notebook(outline):
 
 
 class GenerateChatHandler(BaseChatHandler):
-    id = "generate-chat-handler"
+    id = "generate"
     name = "Generate Notebook"
-    description = "Generates a Jupyter notebook given a description"
     help = "Generates a Jupyter notebook, including name, outline, and section contents"
+    routing_method = "slash_command"
     slash_id = "generate"
 
     def __init__(self, root_dir: str, *args, **kwargs):

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

@@ -29,10 +29,10 @@ def HelpMessage():
 
 
 class HelpChatHandler(BaseChatHandler):
-    id = "help-chat-handler"
+    id = "help"
     name = "Help"
-    description = "Displays information about available commands"
     help = "Displays a help message in the chat message area"
+    routing_method = "slash_command"
     slash_id = "help"
 
     def __init__(self, *args, **kwargs):

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

@@ -31,10 +31,10 @@
 
 
 class LearnChatHandler(BaseChatHandler):
-    id = "learn-chat-handler"
+    id = "learn"
     name = "Learn Local Data"
-    description = "Embed a list of files and directories for use with /ask"
     help = "Pass a list of files and directories. Once converted to vector format, you can ask about them with /ask."
+    routing_method = "slash_command"
     slash_id = "learn"
 
     def __init__(

packages/jupyter-ai/jupyter_ai/extension.py

@@ -42,15 +42,15 @@ class AiExtension(ExtensionApp):
     allowed_providers = List(
         Unicode(),
         default_value=None,
-        help="Identifiers of allow-listed providers. If `None`, all are allowed.",
+        help="Identifiers of allowlisted providers. If `None`, all are allowed.",
         allow_none=True,
         config=True,
     )
 
     blocked_providers = List(
         Unicode(),
         default_value=None,
-        help="Identifiers of block-listed providers. If `None`, none are blocked.",
+        help="Identifiers of blocklisted providers. If `None`, none are blocked.",
         allow_none=True,
         config=True,
     )
@@ -109,19 +109,38 @@ def initialize_settings(self):
         self.log.info("Found entry point groups " + ", ".join(sorted(eps.groups)))
         chat_handler_eps = eps.select(group="jupyter_ai.chat_handlers")
 
-        jai_chat_handlers = {}
-
         chat_handler_kwargs = {
             "log": self.log,
             "config_manager": self.settings["jai_config_manager"],
             "root_chat_handlers": self.settings["jai_root_chat_handlers"],
-            # Everything below this line is the union of all arguments used to
-            # instantiate chat handlers.
-            "chat_history": self.settings["chat_history"],
-            "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,
+        }
+
+        default_chat_handler = DefaultChatHandler(
+            **chat_handler_kwargs, chat_history=self.settings["chat_history"]
+        )
+        clear_chat_handler = ClearChatHandler(
+            **chat_handler_kwargs, chat_history=self.settings["chat_history"]
+        )
+        generate_chat_handler = GenerateChatHandler(
+            **chat_handler_kwargs,
+            root_dir=self.serverapp.root_dir,
+        )
+        learn_chat_handler = LearnChatHandler(
+            **chat_handler_kwargs,
+            root_dir=self.serverapp.root_dir,
+            dask_client_future=dask_client_future,
+        )
+        help_chat_handler = HelpChatHandler(**chat_handler_kwargs)
+        retriever = Retriever(learn_chat_handler=learn_chat_handler)
+        ask_chat_handler = AskChatHandler(**chat_handler_kwargs, retriever=retriever)
+
+        jai_chat_handlers = {
+            "default": default_chat_handler,
+            "/ask": ask_chat_handler,
+            "/clear": clear_chat_handler,
+            "/generate": generate_chat_handler,
+            "/learn": learn_chat_handler,
+            "/help": help_chat_handler,
         }
 
         for chat_handler_ep in chat_handler_eps:
@@ -160,10 +179,8 @@ def initialize_settings(self):
 
         self.settings["jai_chat_handlers"] = jai_chat_handlers
 
-        """
         retriever = Retriever(learn_chat_handler=learn_chat_handler)
         ask_chat_handler = AskChatHandler(**chat_handler_kwargs, retriever=retriever)
-        """
         latency_ms = round((time.time() - start) * 1000)
         self.log.info(f"Initialized Jupyter AI server extension in {latency_ms} ms.")
 

packages/jupyter-ai/pyproject.toml

@@ -41,12 +41,7 @@ dependencies = [
 dynamic = ["version", "description", "authors", "urls", "keywords"]
 
 [project.entry-points."jupyter_ai.chat_handlers"]
-ask = "jupyter_ai:chat_handlers:AskChatHandler"
-clear = "jupyter_ai:chat_handlers:ClearChatHandler"
-default = "jupyter_ai:chat_handlers:DefaultChatHandler"
-generate = "jupyter_ai:chat_handlers:GenerateChatHandler"
-help = "jupyter_ai:chat_handlers:HelpChatHandler"
-learn = "jupyter_ai:chat_handlers:LearnChatHandler"
+custom = "jupyter_ai:chat_handlers.CustomChatHandler"
 
 [project.entry-points."jupyter_ai.default_tasks"]
 core_default_tasks = "jupyter_ai:tasks"