Backport PR #1257 on branch 2.x (Refactor Chat Handlers to Simplify Initialization) (#1266)

* Refactor Chat Handlers to Simplify Initialization (#1257)

* simplify-entrypoints-loading

* fix-lint

* fix-tests

* add-retriever-typing

* remove-retriever-from-base

* fix-circular-import(ydoc-import)

* fix-tests

* fix-type-check-failure

* refactor-retriever-init

(cherry picked from commit 055d7b8)

* fix-test-backporting

* fix-ydoc-unwanted-import

* fix-test

* remove-exception

* reorder-chat-handlers-storage

Author: Darshan808

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

@@ -8,6 +8,7 @@
 from langchain_core.prompts import PromptTemplate
 
 from .base import BaseChatHandler, SlashCommandRoutingType
+from .learn import LearnChatHandler, Retriever
 
 PROMPT_TEMPLATE = """Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question.
 
@@ -18,6 +19,16 @@
 CONDENSE_PROMPT = PromptTemplate.from_template(PROMPT_TEMPLATE)
 
 
+class CustomLearnException(Exception):
+    """Exception raised when Jupyter AI's /ask command is used without the required /learn command."""
+
+    def __init__(self):
+        super().__init__(
+            "Jupyter AI's default /ask command requires the default /learn command. "
+            "If you are overriding /learn via the entry points API, be sure to also override or disable /ask."
+        )
+
+
 class AskChatHandler(BaseChatHandler):
     """Processes messages prefixed with /ask. This actor will
     send the message as input to a RetrieverQA chain, that
@@ -33,12 +44,16 @@ class AskChatHandler(BaseChatHandler):
 
     uses_llm = True
 
-    def __init__(self, retriever, *args, **kwargs):
+    def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        self._retriever = retriever
         self.parser.prog = "/ask"
         self.parser.add_argument("query", nargs=argparse.REMAINDER)
+        learn_chat_handler = self.chat_handlers.get("/learn")
+        if not isinstance(learn_chat_handler, LearnChatHandler):
+            raise CustomLearnException()
+
+        self._retriever = Retriever(learn_chat_handler=learn_chat_handler)
 
     def create_llm_chain(
         self, provider: Type[BaseProvider], provider_params: Dict[str, str]
@@ -51,6 +66,7 @@ def create_llm_chain(
         memory = ConversationBufferWindowMemory(
             memory_key="chat_history", return_messages=True, k=2
         )
+
         self.llm_chain = ConversationalRetrievalChain.from_llm(
             self.llm,
             self._retriever,

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

@@ -4,6 +4,7 @@
 import os
 import time
 import traceback
+from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -156,6 +157,7 @@ def __init__(
         chat_handlers: Dict[str, "BaseChatHandler"],
         context_providers: Dict[str, "BaseCommandContextProvider"],
         message_interrupted: Dict[str, asyncio.Event],
+        log_dir: Optional[str],
     ):
         self.log = log
         self.config_manager = config_manager
@@ -178,6 +180,7 @@ def __init__(
         self.chat_handlers = chat_handlers
         self.context_providers = context_providers
         self.message_interrupted = message_interrupted
+        self.log_dir = Path(log_dir) if log_dir else None
 
         self.llm: Optional[BaseProvider] = None
         self.llm_params: Optional[dict] = None

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

@@ -253,9 +253,8 @@ class GenerateChatHandler(BaseChatHandler):
 
     uses_llm = True
 
-    def __init__(self, log_dir: Optional[str], *args, **kwargs):
+    def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.log_dir = Path(log_dir) if log_dir else None
         self.llm: Optional[BaseProvider] = None
 
     def create_llm_chain(

packages/jupyter-ai/jupyter_ai/extension.py

@@ -5,7 +5,6 @@
 
 from dask.distributed import Client as DaskClient
 from importlib_metadata import entry_points
-from jupyter_ai.chat_handlers.learn import Retriever
 from jupyter_ai_magics import BaseProvider, JupyternautPersona
 from jupyter_ai_magics.utils import get_em_providers, get_lm_providers
 from jupyter_server.extension.application import ExtensionApp
@@ -367,16 +366,17 @@ def _init_chat_handlers(self):
             "chat_handlers": chat_handlers,
             "context_providers": self.settings["jai_context_providers"],
             "message_interrupted": self.settings["jai_message_interrupted"],
+            "log_dir": self.error_logs_dir,
         }
+
         default_chat_handler = DefaultChatHandler(**chat_handler_kwargs)
+        generate_chat_handler = GenerateChatHandler(**chat_handler_kwargs)
         clear_chat_handler = ClearChatHandler(**chat_handler_kwargs)
-        generate_chat_handler = GenerateChatHandler(
-            **chat_handler_kwargs,
-            log_dir=self.error_logs_dir,
-        )
         learn_chat_handler = LearnChatHandler(**chat_handler_kwargs)
-        retriever = Retriever(learn_chat_handler=learn_chat_handler)
-        ask_chat_handler = AskChatHandler(**chat_handler_kwargs, retriever=retriever)
+        # Store learn_chat_handler before initializing AskChatHandler,
+        # as it is required for initializing the Retriever.
+        chat_handlers["/learn"] = learn_chat_handler
+        ask_chat_handler = AskChatHandler(**chat_handler_kwargs)
 
         export_chat_handler = ExportChatHandler(**chat_handler_kwargs)
 
@@ -386,7 +386,6 @@ def _init_chat_handlers(self):
         chat_handlers["/ask"] = ask_chat_handler
         chat_handlers["/clear"] = clear_chat_handler
         chat_handlers["/generate"] = generate_chat_handler
-        chat_handlers["/learn"] = learn_chat_handler
         chat_handlers["/export"] = export_chat_handler
         chat_handlers["/fix"] = fix_chat_handler
 

packages/jupyter-ai/jupyter_ai/tests/test_handlers.py

@@ -78,6 +78,7 @@ def broadcast_message(message: Message) -> None:
             chat_handlers={},
             context_providers={},
             message_interrupted={},
+            log_dir="",
         )