Introduce AI persona framework (#1324)

* add dev-reinstall script for convenience

* add PersonaManager and JupyternautPersona

* add PersonaAwareness to allow awareness on >1 persona

* add DebugPersona to jupyter_ai_test

* upgrade to Jupyter Chat v0.10.0

* automatically generate IDs for personas

* add mention-based routing to PersonaManager

* implement streaming replies in Jupyternaut persona

* correctly identify AI messages in YChatHistory

* fix mention routing to actually work

* add logging capability to PersonaAwareness

* update BasePersona.id to not depend on module path

* improve logging in PersonaManager

* log time elapsed in PersonaManager

* simplify jupyternaut prompt template

* pre-commit

* add comment to ref jupyter-chat#212

* have _init_persona_manager() explicitly return None on exception

* pre-commit

* update docstrings on BasePersona

* pre-commit

* add return type annotation to forward_reply_stream()

* remove unused return

* add comment explaining _init_persona_classes()

* add docstrings to persona_manager methods

* remove redundant error logs

* pre-commit

* fix mypy errors

* remove tests on chat handlers as they are superseded by personas

* fix type annotation for Py39 compat

* rename forward_reply_stream() to stream_message()

* add send_message() method to BasePersona

* pre-commit

Author: dlqqq

package.json

@@ -30,6 +30,7 @@
     "clean:all": "lerna run clean:all",
     "dev": "jupyter lab --config playground/config.py",
     "dev-install": "lerna run dev-install --stream",
+    "dev-reinstall": "jlpm dev-uninstall && jlpm dev-install",
     "dev-uninstall": "lerna run dev-uninstall --stream",
     "install-from-src": "lerna run install-from-src --stream",
     "lint": "jlpm && lerna run prettier && lerna run eslint",

packages/jupyter-ai-test/jupyter_ai_test/debug_persona.py

@@ -0,0 +1,24 @@
+from jupyter_ai.personas.base_persona import BasePersona, PersonaDefaults
+from jupyterlab_chat.models import Message, NewMessage
+
+
+class DebugPersona(BasePersona):
+    """
+    The Jupyternaut persona, the main persona provided by Jupyter AI.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @property
+    def defaults(self):
+        return PersonaDefaults(
+            name="DebugPersona",
+            avatar_path="/api/ai/static/jupyternaut.svg",
+            description="A mock persona used for debugging in local dev environments.",
+            system_prompt="...",
+        )
+
+    async def process_message(self, message: Message):
+        self.ychat.add_message(NewMessage(body="Hello!", sender=self.id))
+        return

packages/jupyter-ai-test/pyproject.toml

@@ -35,6 +35,9 @@ test-provider-ask-learn-unsupported = "jupyter_ai_test.test_providers:TestProvid
 [project.entry-points."jupyter_ai.chat_handlers"]
 test-slash-command = "jupyter_ai_test.test_slash_commands:TestSlashCommand"
 
+[project.entry-points."jupyter_ai.personas"]
+debug-persona = "jupyter_ai_test.debug_persona:DebugPersona"
+
 [tool.hatch.build.hooks.version]
 path = "jupyter_ai_test/_version.py"
 

packages/jupyter-ai/jupyter_ai/extension.py

@@ -2,8 +2,9 @@
 import re
 import time
 import types
+from asyncio import get_event_loop_policy
 from functools import partial
-from typing import Dict
+from typing import TYPE_CHECKING, Dict, Optional
 
 import traitlets
 from dask.distributed import Client as DaskClient
@@ -12,7 +13,6 @@
 from jupyter_ai_magics.utils import get_em_providers, get_lm_providers
 from jupyter_events import EventLogger
 from jupyter_server.extension.application import ExtensionApp
-from jupyter_server.utils import url_path_join
 from jupyterlab_chat.models import Message
 from jupyterlab_chat.ychat import YChat
 from pycrdt import ArrayEvent
@@ -22,7 +22,6 @@
 from .chat_handlers.base import BaseChatHandler
 from .completions.handlers import DefaultInlineCompletionHandler
 from .config_manager import ConfigManager
-from .constants import BOT
 from .context_providers import BaseCommandContextProvider, FileContextProvider
 from .handlers import (
     ApiKeysHandler,
@@ -33,6 +32,10 @@
     SlashCommandsInfoHandler,
 )
 from .history import YChatHistory
+from .personas import PersonaManager
+
+if TYPE_CHECKING:
+    from asyncio import AbstractEventLoop
 
 from jupyter_collaboration import (  # type:ignore[import-untyped]  # isort:skip
     __version__ as jupyter_collaboration_version,
@@ -244,6 +247,13 @@ def initialize(self):
             schema_id=JUPYTER_COLLABORATION_EVENTS_URI, listener=self.connect_chat
         )
 
+    @property
+    def event_loop(self) -> "AbstractEventLoop":
+        """
+        Returns a reference to the asyncio event loop.
+        """
+        return get_event_loop_policy().get_event_loop()
+
     async def connect_chat(
         self, logger: EventLogger, schema_id: str, data: dict
     ) -> None:
@@ -264,17 +274,19 @@ async def connect_chat(
         if ychat is None:
             return
 
-        # Add the bot user to the chat document awareness.
-        BOT["avatar_url"] = url_path_join(
-            self.settings.get("base_url", "/"), "api/ai/static/jupyternaut.svg"
-        )
-        if ychat.awareness is not None:
-            ychat.awareness.set_local_state_field("user", BOT)
-
         # initialize chat handlers for new chat
         self.chat_handlers_by_room[room_id] = self._init_chat_handlers(ychat)
 
-        callback = partial(self.on_change, room_id)
+        # initialize persona manager
+        persona_manager = self._init_persona_manager(ychat)
+        if not persona_manager:
+            self.log.error(
+                "Jupyter AI was unable to initialize its AI personas. They are not available for use in chat until this error is resolved. "
+                + "Please verify your configuration and open a new issue on GitHub if this error persists."
+            )
+            return
+
+        callback = partial(self.on_change, room_id, persona_manager)
         ychat.ymessages.observe(callback)
 
     async def get_chat(self, room_id: str) -> YChat:
@@ -301,21 +313,26 @@ async def get_chat(self, room_id: str) -> YChat:
         self.ychats_by_room[room_id] = document
         return document
 
-    def on_change(self, room_id: str, events: ArrayEvent) -> None:
+    def on_change(
+        self, room_id: str, persona_manager: PersonaManager, events: ArrayEvent
+    ) -> None:
         assert self.serverapp
 
         for change in events.delta:  # type:ignore[attr-defined]
             if not "insert" in change.keys():
                 continue
-            messages = change["insert"]
-            for message_dict in messages:
-                message = Message(**message_dict)
-                if message.sender == BOT["username"] or message.raw_time:
-                    continue
 
-                self.serverapp.io_loop.asyncio_loop.create_task(  # type:ignore[attr-defined]
-                    self.route_human_message(room_id, message)
-                )
+            # the "if not m['raw_time']" clause is necessary because every new
+            # message triggers 2 events, one with `raw_time` set to `True` and
+            # another with `raw_time` set to `False` milliseconds later.
+            # we should explore fixing this quirk in Jupyter Chat.
+            #
+            # Ref: https://github.com/jupyterlab/jupyter-chat/issues/212
+            new_messages = [
+                Message(**m) for m in change["insert"] if not m.get("raw_time", False)
+            ]
+            for new_message in new_messages:
+                persona_manager.route_message(new_message)
 
     async def route_human_message(self, room_id: str, message: Message):
         """
@@ -400,18 +417,15 @@ def initialize_settings(self):
 
         self.log.info(f"Registered {self.name} server extension")
 
-        # get reference to event loop
-        # `asyncio.get_event_loop()` is deprecated in Python 3.11+, in favor of
-        # the more readable `asyncio.get_event_loop_policy().get_event_loop()`.
-        # it's easier to just reference the loop directly.
-        loop = self.serverapp.io_loop.asyncio_loop
-        self.settings["jai_event_loop"] = loop
+        self.settings["jai_event_loop"] = self.event_loop
 
         # We cannot instantiate the Dask client directly here because it
         # requires the event loop to be running on init. So instead we schedule
         # this as a task that is run as soon as the loop starts, and pass
         # consumers a Future that resolves to the Dask client when awaited.
-        self.settings["dask_client_future"] = loop.create_task(self._get_dask_client())
+        self.settings["dask_client_future"] = self.event_loop.create_task(
+            self._get_dask_client()
+        )
 
         # Create empty context providers dict to be filled later.
         # This is created early to use as kwargs for chat handlers.
@@ -456,10 +470,7 @@ async def _stop_extension(self):
 
     def _init_chat_handlers(self, ychat: YChat) -> Dict[str, BaseChatHandler]:
         """
-        Initializes a set of chat handlers. May accept a YChat instance for
-        collaborative chats.
-
-        TODO: Make `ychat` required once Jupyter Chat migration is complete.
+        Initializes a set of chat handlers for a given `YChat` instance.
         """
         assert self.serverapp
 
@@ -606,3 +617,32 @@ def _init_context_providers(self):
                 **context_providers_kwargs
             )
             self.log.info(f"Registered context provider `{context_provider.id}`.")
+
+    def _init_persona_manager(self, ychat: YChat) -> Optional[PersonaManager]:
+        """
+        Initializes a `PersonaManager` instance scoped to a `YChat`.
+
+        This method should not raise an exception. Upon encountering an
+        exception, this method will catch it, log it, and return `None`.
+        """
+        persona_manager: Optional[PersonaManager]
+
+        try:
+            config_manager = self.settings.get("jai_config_manager", None)
+            assert config_manager and isinstance(config_manager, ConfigManager)
+
+            persona_manager = PersonaManager(
+                ychat=ychat,
+                config_manager=config_manager,
+                event_loop=self.event_loop,
+                log=self.log,
+            )
+        except Exception as e:
+            # TODO: how to stop the extension when this fails
+            # also why do uncaught exceptions produce an empty error log in Jupyter Server?
+            self.log.error(
+                f"Unable to initialize PersonaManager in YChat with ID '{ychat.get_id()}' due to an exception printed below."
+            )
+            self.log.exception(e)
+        finally:
+            return persona_manager

packages/jupyter-ai/jupyter_ai/history.py

@@ -1,6 +1,5 @@
 from typing import List, Optional
 
-from jupyter_ai.constants import BOT
 from jupyterlab_chat.models import Message as JChatMessage
 from jupyterlab_chat.ychat import YChat
 from langchain_core.chat_history import BaseChatMessageHistory
@@ -46,7 +45,7 @@ def _convert_to_langchain_messages(self, jchat_messages: List[JChatMessage]):
         """
         messages: List[BaseMessage] = []
         for jchat_message in jchat_messages:
-            if jchat_message.sender == BOT["username"]:
+            if jchat_message.sender.startswith("jupyter-ai-personas::"):
                 messages.append(AIMessage(content=jchat_message.body))
             else:
                 messages.append(HumanMessage(content=jchat_message.body))

packages/jupyter-ai/jupyter_ai/personas/__init__.py

@@ -0,0 +1,2 @@
+from .base_persona import BasePersona, PersonaDefaults
+from .persona_manager import PersonaManager