From 4abacfbeb5c72e365ec20351f8579d92dbf83437 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Krassowski?=
 <5832902+krassowski@users.noreply.github.com>
Date: Tue, 19 Dec 2023 01:14:34 +0000
Subject: [PATCH 1/2] Dynamically generate help message for slash commands in
 chat UI (#520)

* Show help for custom slash commands

* Morph `short_help` into existing `help` per review
---
 .../jupyter_ai/chat_handlers/ask.py           |  2 +-
 .../jupyter_ai/chat_handlers/clear.py         |  2 +-
 .../jupyter_ai/chat_handlers/generate.py      |  2 +-
 .../jupyter_ai/chat_handlers/help.py          | 30 ++++++++++++-------
 .../jupyter_ai/chat_handlers/learn.py         |  2 +-
 packages/jupyter-ai/jupyter_ai/extension.py   | 14 +++++++--
 6 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py
index bfb55ce21..5716ddd9b 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py
@@ -28,7 +28,7 @@ class AskChatHandler(BaseChatHandler):
 
     id = "ask"
     name = "Ask with Local Data"
-    help = "Asks a question with retrieval augmented generation (RAG)"
+    help = "Ask a question about your learned data"
     routing_type = SlashCommandRoutingType(slash_id="ask")
 
     def __init__(self, retriever, *args, **kwargs):
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py
index 7042c4632..73db24a80 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py
@@ -8,7 +8,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"
+    help = "Clear the chat window"
     routing_type = SlashCommandRoutingType(slash_id="clear")
 
     def __init__(self, *args, **kwargs):
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py
index a17f51390..1f463dd9e 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py
@@ -218,7 +218,7 @@ def create_notebook(outline):
 class GenerateChatHandler(BaseChatHandler):
     id = "generate"
     name = "Generate Notebook"
-    help = "Generates a Jupyter notebook, including name, outline, and section contents"
+    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):
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py
index cbf4c19c9..3735a7c5e 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py
@@ -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
@@ -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.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="",
     )
 
@@ -31,11 +38,12 @@ def HelpMessage():
 class HelpChatHandler(BaseChatHandler):
     id = "help"
     name = "Help"
-    help = "Displays a help message in the chat message area"
+    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)
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py
index 68a59bef1..a36885297 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py
@@ -33,7 +33,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."
+    help = "Teach Jupyternaut about files on your system"
     routing_type = SlashCommandRoutingType(slash_id="learn")
 
     def __init__(self, *args, **kwargs):
diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py
index c6ce4f5d9..a41c556c7 100644
--- a/packages/jupyter-ai/jupyter_ai/extension.py
+++ b/packages/jupyter-ai/jupyter_ai/extension.py
@@ -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
@@ -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)
 
@@ -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:
@@ -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)

From 79e345a57db9aa170727a254410783cc422a6654 Mon Sep 17 00:00:00 2001
From: david qiu <david@qiu.dev>
Date: Tue, 19 Dec 2023 11:02:16 -0800
Subject: [PATCH 2/2] Handle LLMs lacking concurrency support in Chat UI (#506)

* specify no concurrency in GPT4All models

* handle LLMs lacking concurrency support in chat

* add annotation & default for uses_llm attr

* edit reply msg
---
 .../jupyter_ai_magics/providers.py            |  5 +++
 .../jupyter_ai/chat_handlers/ask.py           |  2 ++
 .../jupyter_ai/chat_handlers/base.py          | 34 ++++++++++++++++---
 .../jupyter_ai/chat_handlers/clear.py         |  2 ++
 .../jupyter_ai/chat_handlers/default.py       |  2 ++
 .../jupyter_ai/chat_handlers/generate.py      |  2 ++
 .../jupyter_ai/chat_handlers/help.py          |  2 ++
 .../jupyter_ai/chat_handlers/learn.py         |  2 ++
 8 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
index 33d220f94..9d843fe4d 100644
--- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
+++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -362,6 +362,11 @@ def __init__(self, **kwargs):
     async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]:
         return await self._call_in_executor(*args, **kwargs)
 
+    @property
+    def allows_concurrency(self):
+        # At present, GPT4All providers fail with concurrent messages. See #481.
+        return False
+
 
 HUGGINGFACE_HUB_VALID_TASKS = (
     "text2text-generation",
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py
index 5716ddd9b..8fe5c7f3a 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py
@@ -31,6 +31,8 @@ class AskChatHandler(BaseChatHandler):
     help = "Ask a question about your learned data"
     routing_type = SlashCommandRoutingType(slash_id="ask")
 
+    uses_llm = True
+
     def __init__(self, retriever, *args, **kwargs):
         super().__init__(*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 15aef6788..fb8a5e3de 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py
@@ -58,6 +58,15 @@ class BaseChatHandler:
 
     routing_type: HandlerRoutingType = ...
 
+    uses_llm: ClassVar[bool] = True
+    """Class attribute specifying whether this chat handler uses the LLM
+    specified by the config. Subclasses should define this. Should be set to
+    `False` for handlers like `/help`."""
+
+    _requests_count = 0
+    """Class attribute set to the number of requests that Jupyternaut is
+    currently handling."""
+
     def __init__(
         self,
         log: Logger,
@@ -82,11 +91,26 @@ def __init__(
 
     async def on_message(self, message: HumanChatMessage):
         """
-        Method which receives a human message and processes it via
-        `self.process_message()`, calling `self.handle_exc()` when an exception
-        is raised. This method is called by RootChatHandler when it routes a
-        human message to this chat handler.
+        Method which receives a human message, calls `self.get_llm_chain()`, and
+        processes the message via `self.process_message()`, calling
+        `self.handle_exc()` when an exception is raised. This method is called
+        by RootChatHandler when it routes a human message to this chat handler.
         """
+
+        # check whether the configured LLM can support a request at this time.
+        if self.uses_llm and BaseChatHandler._requests_count > 0:
+            lm_provider_klass = self.config_manager.lm_provider
+            lm_provider_params = self.config_manager.lm_provider_params
+            lm_provider = lm_provider_klass(**lm_provider_params)
+
+            if not lm_provider.allows_concurrency:
+                self.reply(
+                    "The currently selected language model can process only one request at a time. Please wait for me to reply before sending another question.",
+                    message,
+                )
+                return
+
+        BaseChatHandler._requests_count += 1
         try:
             await self.process_message(message)
         except Exception as e:
@@ -96,6 +120,8 @@ async def on_message(self, message: HumanChatMessage):
                 await self.handle_exc(e, message)
             except Exception as e:
                 await self._default_handle_exc(e, message)
+        finally:
+            BaseChatHandler._requests_count -= 1
 
     async def process_message(self, message: HumanChatMessage):
         """
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py
index 73db24a80..0555abfca 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py
@@ -11,6 +11,8 @@ class ClearChatHandler(BaseChatHandler):
     help = "Clear the chat window"
     routing_type = SlashCommandRoutingType(slash_id="clear")
 
+    uses_llm = False
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py
index 5bd839ca5..3a76fba44 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py
@@ -37,6 +37,8 @@ class DefaultChatHandler(BaseChatHandler):
     help = "Responds to prompts that are not otherwise handled by a chat handler"
     routing_type = SlashCommandRoutingType(slash_id=None)
 
+    uses_llm = True
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.memory = ConversationBufferWindowMemory(return_messages=True, k=2)
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py
index 1f463dd9e..50ef6b02c 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py
@@ -221,6 +221,8 @@ class GenerateChatHandler(BaseChatHandler):
     help = "Generate a Jupyter notebook from a text prompt"
     routing_type = SlashCommandRoutingType(slash_id="generate")
 
+    uses_llm = True
+
     def __init__(self, preferred_dir: str, log_dir: Optional[str], *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.log_dir = Path(log_dir) if log_dir else None
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py
index 3735a7c5e..ebb8f0383 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py
@@ -41,6 +41,8 @@ class HelpChatHandler(BaseChatHandler):
     help = "Display this help message"
     routing_type = SlashCommandRoutingType(slash_id="help")
 
+    uses_llm = False
+
     def __init__(self, *args, chat_handlers: Dict[str, BaseChatHandler], **kwargs):
         super().__init__(*args, **kwargs)
         self._chat_handlers = chat_handlers
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py
index a36885297..3f67c07d2 100644
--- a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py
+++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py
@@ -36,6 +36,8 @@ class LearnChatHandler(BaseChatHandler):
     help = "Teach Jupyternaut about files on your system"
     routing_type = SlashCommandRoutingType(slash_id="learn")
 
+    uses_llm = True
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.parser.prog = "/learn"