Skip to content

Commit

Permalink
Show help for custom slash commands
Browse files Browse the repository at this point in the history
  • Loading branch information
krassowski authored and JasonWeill committed Dec 18, 2023
1 parent 54daef8 commit de0ca5e
Show file tree
Hide file tree
Showing 8 changed files with 38 additions and 13 deletions.
1 change: 1 addition & 0 deletions docs/source/developers/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,7 @@ class CustomChatHandler(BaseChatHandler):
id = "custom"
name = "Custom"
help = "A chat handler that does something custom"
short_help = "Do something custom"
routing_type = SlashCommandRoutingType(slash_id="custom")

def __init__(self, *args, **kwargs):
Expand Down
1 change: 1 addition & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ class AskChatHandler(BaseChatHandler):
id = "ask"
name = "Ask with Local Data"
help = "Asks a question with retrieval augmented generation (RAG)"
short_help = "Ask a question about your learned data"
routing_type = SlashCommandRoutingType(slash_id="ask")

def __init__(self, retriever, *args, **kwargs):
Expand Down
3 changes: 3 additions & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,9 @@ class BaseChatHandler:
"""What this chat handler does, which third-party models it contacts,
the data it returns to the user, and so on, for display in the UI."""

short_help: ClassVar[Optional[str]] = None
"""A shorter version of help message for display in the chat UI."""

routing_type: HandlerRoutingType = ...

def __init__(
Expand Down
1 change: 1 addition & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ class ClearChatHandler(BaseChatHandler):
id = "clear"
name = "Clear chat messages"
help = "Clears the displayed chat message history only; does not clear the context sent to chat providers"
short_help = "Clear the chat window"
routing_type = SlashCommandRoutingType(slash_id="clear")

def __init__(self, *args, **kwargs):
Expand Down
1 change: 1 addition & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,7 @@ class GenerateChatHandler(BaseChatHandler):
id = "generate"
name = "Generate Notebook"
help = "Generates a Jupyter notebook, including name, outline, and section contents"
short_help = "Generate a Jupyter notebook from a text prompt"
routing_type = SlashCommandRoutingType(slash_id="generate")

def __init__(self, preferred_dir: str, log_dir: Optional[str], *args, **kwargs):
Expand Down
29 changes: 19 additions & 10 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/help.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import time
from typing import List
from typing import Dict
from uuid import uuid4

from jupyter_ai.models import AgentChatMessage, HumanChatMessage
Expand All @@ -8,22 +8,29 @@

HELP_MESSAGE = """Hi there! I'm Jupyternaut, your programming assistant.
You can ask me a question using the text box below. You can also use these commands:
* `/learn` — Teach Jupyternaut about files on your system
* `/ask` — Ask a question about your learned data
* `/generate` — Generate a Jupyter notebook from a text prompt
* `/clear` — Clear the chat window
* `/help` — Display this help message
{commands}
Jupyter AI includes [magic commands](https://jupyter-ai.readthedocs.io/en/latest/users/index.html#the-ai-and-ai-magic-commands) that you can use in your notebooks.
For more information, see the [documentation](https://jupyter-ai.readthedocs.io).
"""


def HelpMessage():
def _format_help_message(chat_handlers: Dict[str, BaseChatHandler]):
commands = "\n".join(
[
f"* `{command_name}` — {handler.short_help or handler.help}"
for command_name, handler in chat_handlers.items()
if command_name != "default"
]
)
return HELP_MESSAGE.format(commands=commands)


def HelpMessage(chat_handlers: Dict[str, BaseChatHandler]):
return AgentChatMessage(
id=uuid4().hex,
time=time.time(),
body=HELP_MESSAGE,
body=_format_help_message(chat_handlers),
reply_to="",
)

Expand All @@ -32,10 +39,12 @@ class HelpChatHandler(BaseChatHandler):
id = "help"
name = "Help"
help = "Displays a help message in the chat message area"
short_help = "Display this help message"
routing_type = SlashCommandRoutingType(slash_id="help")

def __init__(self, *args, **kwargs):
def __init__(self, *args, chat_handlers: Dict[str, BaseChatHandler], **kwargs):
super().__init__(*args, **kwargs)
self._chat_handlers = chat_handlers

async def process_message(self, message: HumanChatMessage):
self.reply(HELP_MESSAGE, message)
self.reply(_format_help_message(self._chat_handlers), message)
1 change: 1 addition & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ class LearnChatHandler(BaseChatHandler):
id = "learn"
name = "Learn Local Data"
help = "Pass a list of files and directories. Once converted to vector format, you can ask about them with /ask."
short_help = "Teach Jupyternaut about files on your system"
routing_type = SlashCommandRoutingType(slash_id="learn")

def __init__(self, *args, **kwargs):
Expand Down
14 changes: 11 additions & 3 deletions packages/jupyter-ai/jupyter_ai/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,7 @@ def initialize_settings(self):
# list of chat messages to broadcast to new clients
# this is only used to render the UI, and is not the conversational
# memory object used by the LM chain.
self.settings["chat_history"] = [HelpMessage()]
self.settings["chat_history"] = []

# get reference to event loop
# `asyncio.get_event_loop()` is deprecated in Python 3.11+, in favor of
Expand Down Expand Up @@ -189,7 +189,6 @@ def initialize_settings(self):
log_dir=self.error_logs_dir,
)
learn_chat_handler = LearnChatHandler(**chat_handler_kwargs)
help_chat_handler = HelpChatHandler(**chat_handler_kwargs)
retriever = Retriever(learn_chat_handler=learn_chat_handler)
ask_chat_handler = AskChatHandler(**chat_handler_kwargs, retriever=retriever)

Expand All @@ -199,9 +198,12 @@ def initialize_settings(self):
"/clear": clear_chat_handler,
"/generate": generate_chat_handler,
"/learn": learn_chat_handler,
"/help": help_chat_handler,
}

help_chat_handler = HelpChatHandler(
**chat_handler_kwargs, chat_handlers=jai_chat_handlers
)

slash_command_pattern = r"^[a-zA-Z0-9_]+$"
for chat_handler_ep in chat_handler_eps:
try:
Expand Down Expand Up @@ -248,6 +250,12 @@ def initialize_settings(self):
f"Registered chat handler `{chat_handler.id}` with command `{command_name}`."
)

# Make help always appear as the last command
jai_chat_handlers["/help"] = help_chat_handler

self.settings["chat_history"].append(
HelpMessage(chat_handlers=jai_chat_handlers)
)
self.settings["jai_chat_handlers"] = jai_chat_handlers

latency_ms = round((time.time() - start) * 1000)
Expand Down

0 comments on commit de0ca5e

Please sign in to comment.