Introduce AI personas (see PR #1324) (#1341)

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