Skip to content

Commit

Permalink
Backport PR jupyterlab#506: Handle LLMs lacking concurrency support i…
Browse files Browse the repository at this point in the history
…n Chat UI
  • Loading branch information
dlqqq authored and meeseeksmachine committed Dec 19, 2023
1 parent 4d2dc86 commit 1350b61
Show file tree
Hide file tree
Showing 8 changed files with 47 additions and 4 deletions.
5 changes: 5 additions & 0 deletions packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
2 changes: 2 additions & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
34 changes: 30 additions & 4 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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:
Expand All @@ -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):
"""
Expand Down
2 changes: 2 additions & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
2 changes: 2 additions & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/default.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/help.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down

0 comments on commit 1350b61

Please sign in to comment.