From b221758e83290d2cbc7ee94450675f782b5e75ed Mon Sep 17 00:00:00 2001 From: Jason Weill Date: Tue, 24 Oct 2023 15:49:08 -0700 Subject: [PATCH] WIP: Updates per feedback, adds custom handler --- .../jupyter_ai/chat_handlers/__init__.py | 1 + .../jupyter_ai/chat_handlers/ask.py | 4 +- .../jupyter_ai/chat_handlers/base.py | 23 +++++----- .../jupyter_ai/chat_handlers/clear.py | 4 +- .../jupyter_ai/chat_handlers/custom.py | 33 ++++++++++++++ .../jupyter_ai/chat_handlers/generate.py | 4 +- .../jupyter_ai/chat_handlers/help.py | 4 +- .../jupyter_ai/chat_handlers/learn.py | 4 +- packages/jupyter-ai/jupyter_ai/extension.py | 43 +++++++++++++------ packages/jupyter-ai/pyproject.toml | 7 +-- 10 files changed, 88 insertions(+), 39 deletions(-) create mode 100644 packages/jupyter-ai/jupyter_ai/chat_handlers/custom.py diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py index c3c64b789..5da53d8ff 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/__init__.py +++ b/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 diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py index 717e20902..21557c57d 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py +++ b/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): diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py index fe0daada5..ba0314c08 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py +++ b/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, diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py index 365403d3e..32a326041 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py +++ b/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): diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/custom.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/custom.py new file mode 100644 index 000000000..2f419721a --- /dev/null +++ b/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) diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py index a49f513cf..f358d8c97 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py +++ b/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): diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py index 474b2b2ab..c167d7bdb 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py +++ b/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): diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py index 64f950c33..db0368f03 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py +++ b/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__( diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index e37580d05..65439d656 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -42,7 +42,7 @@ 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, ) @@ -50,7 +50,7 @@ class AiExtension(ExtensionApp): 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.") diff --git a/packages/jupyter-ai/pyproject.toml b/packages/jupyter-ai/pyproject.toml index 63ce38cb9..e9e822925 100644 --- a/packages/jupyter-ai/pyproject.toml +++ b/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"