Allow personas to get chat path and directory (#1379)

* add warning on v2 persona model

* add get_chat_path() and get_chat_dir() methods to persona manager

* make PersonaManager logging configurable

* make BasePersona logging configurable

* add chat path & dir methods to BasePersona

* pre-commit

* fix mypy errors

* pre-commit

* fix mypy

* make get_chat_path() absolute by default

* fixup previous commit

Author: dlqqq

packages/jupyter-ai-magics/jupyter_ai_magics/models/persona.py

@@ -1,3 +1,11 @@
+# NOTE: This is the outdated `Persona` model used by Jupyter AI v2.
+# This is deprecated and will be removed by Jupyter AI v3.
+# The latest definition of a persona is located in
+# `jupyter_ai/personas/base_persona.py`.
+#
+# TODO: Delete this file once v3 model API changes are complete. The current model
+# API still depends on this, so that work must be done first.
+
 from pydantic import BaseModel
 
 

packages/jupyter-ai/jupyter_ai/extension.py

@@ -10,11 +10,14 @@
 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_fileid.manager import (  # type: ignore[import-untyped]
+    BaseFileIdManager,
+)
 from jupyterlab_chat.models import Message
 from jupyterlab_chat.ychat import YChat
 from pycrdt import ArrayEvent
 from tornado.web import StaticFileHandler
-from traitlets import Integer, List, Unicode
+from traitlets import Integer, List, Type, Unicode
 
 from .completions.handlers import DefaultInlineCompletionHandler
 from .config_manager import ConfigManager
@@ -71,6 +74,13 @@ class AiExtension(ExtensionApp):
         ),
     ]
 
+    persona_manager_class = Type(
+        klass=PersonaManager,
+        default_value=PersonaManager,
+        config=True,
+        help="The `PersonaManager` class.",
+    )
+
     allowed_providers = List(
         Unicode(),
         default_value=None,
@@ -230,7 +240,7 @@ async def connect_chat(
             return
 
         # initialize persona manager
-        persona_manager = self._init_persona_manager(ychat)
+        persona_manager = self._init_persona_manager(room_id, 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. "
@@ -372,7 +382,9 @@ async def _stop_extension(self):
         """
         # TODO: explore if cleanup is necessary
 
-    def _init_persona_manager(self, ychat: YChat) -> Optional[PersonaManager]:
+    def _init_persona_manager(
+        self, room_id: str, ychat: YChat
+    ) -> Optional[PersonaManager]:
         """
         Initializes a `PersonaManager` instance scoped to a `YChat`.
 
@@ -390,11 +402,27 @@ def _init_persona_manager(self, ychat: YChat) -> Optional[PersonaManager]:
                 message_interrupted, dict
             )
 
-            persona_manager = PersonaManager(
+            assert self.serverapp
+            assert self.serverapp.web_app
+            assert self.serverapp.web_app.settings
+            fileid_manager = self.serverapp.web_app.settings.get(
+                "file_id_manager", None
+            )
+            assert isinstance(fileid_manager, BaseFileIdManager)
+
+            contents_manager = self.serverapp.contents_manager
+            root_dir = getattr(contents_manager, "root_dir", None)
+            assert isinstance(root_dir, str)
+
+            PersonaManagerClass = self.persona_manager_class
+            persona_manager = PersonaManagerClass(
+                parent=self,
+                room_id=room_id,
                 ychat=ychat,
                 config_manager=config_manager,
+                fileid_manager=fileid_manager,
+                root_dir=root_dir,
                 event_loop=self.event_loop,
-                log=self.log,
                 message_interrupted=message_interrupted,
             )
         except Exception as e:

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

@@ -1,5 +1,5 @@
 import asyncio
-from abc import ABC, abstractmethod
+from abc import ABC, ABCMeta, abstractmethod
 from dataclasses import asdict
 from logging import Logger
 from time import time
@@ -9,6 +9,8 @@
 from jupyterlab_chat.models import Message, NewMessage, User
 from jupyterlab_chat.ychat import YChat
 from pydantic import BaseModel
+from traitlets import MetaHasTraits
+from traitlets.config import LoggingConfigurable
 
 from .persona_awareness import PersonaAwareness
 
@@ -44,7 +46,15 @@ class PersonaDefaults(BaseModel):
     # ^^^ set this to automatically default to a model after a fresh start, no config file
 
 
-class BasePersona(ABC):
+class ABCLoggingConfigurableMeta(ABCMeta, MetaHasTraits):
+    """
+    Metaclass required for `BasePersona` to inherit from both `ABC` and
+    `LoggingConfigurable`. This pattern is also followed by `BaseFileIdManager`
+    from `jupyter_server_fileid`.
+    """
+
+
+class BasePersona(ABC, LoggingConfigurable, metaclass=ABCLoggingConfigurableMeta):
     """
     Abstract base class that defines a persona when implemented.
     """
@@ -55,21 +65,22 @@ class BasePersona(ABC):
     Automatically set by `BasePersona`.
     """
 
-    manager: "PersonaManager"
+    parent: "PersonaManager"  # type: ignore
     """
     Reference to the `PersonaManager` for this `YChat`, which manages this
-    instance. Automatically set by `BasePersona`.
+    instance. Automatically set by the `LoggingConfigurable` parent class.
     """
 
-    config: ConfigManager
+    config_manager: ConfigManager
     """
     Reference to the `ConfigManager` singleton, which is used to read & write from
     the Jupyter AI settings. Automatically set by `BasePersona`.
     """
 
-    log: Logger
+    log: Logger  # type: ignore
     """
-    The logger for this instance. Automatically set by `BasePersona`.
+    The `logging.Logger` instance used by this class. Automatically set by the
+    `LoggingConfigurable` parent class.
     """
 
     awareness: PersonaAwareness
@@ -92,22 +103,26 @@ class BasePersona(ABC):
     ################################################
     def __init__(
         self,
-        *,
+        *args,
         ychat: YChat,
-        manager: "PersonaManager",
-        config: ConfigManager,
-        log: Logger,
+        config_manager: ConfigManager,
         message_interrupted: dict[str, asyncio.Event],
+        **kwargs,
     ):
+        # Forward other arguments to parent class
+        super().__init__(*args, **kwargs)
+
+        # Bind arguments to instance attributes
         self.ychat = ychat
-        self.manager = manager
-        self.config = config
-        self.log = log
+        self.config_manager = config_manager
         self.message_interrupted = message_interrupted
+
+        # Initialize custom awareness object for this persona
         self.awareness = PersonaAwareness(
             ychat=self.ychat, log=self.log, user=self.as_user()
         )
 
+        # Register this persona as a user in the chat
         self.ychat.set_user(self.as_user())
 
     ################################################
@@ -298,6 +313,22 @@ def send_message(self, body: str) -> None:
         """
         self.ychat.add_message(NewMessage(body=body, sender=self.id))
 
+    def get_chat_path(self, relative: bool = False) -> str:
+        """
+        Returns the absolute path of the chat file assigned to this persona.
+
+        To get a path relative to the `ContentsManager` root directory, call
+        this method with `relative=True`.
+        """
+        return self.parent.get_chat_path(relative=relative)
+
+    def get_chat_dir(self) -> str:
+        """
+        Returns the absolute path to the parent directory of the chat file
+        assigned to this persona.
+        """
+        return self.parent.get_chat_dir()
+
 
 class GenerationInterrupted(asyncio.CancelledError):
     """Exception raised when streaming is cancelled by the user"""

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

@@ -27,8 +27,8 @@ def defaults(self):
         )
 
     async def process_message(self, message: Message) -> None:
-        provider_name = self.config.lm_provider.name
-        model_id = self.config.lm_provider_params["model_id"]
+        provider_name = self.config_manager.lm_provider.name
+        model_id = self.config_manager.lm_provider_params["model_id"]
 
         runnable = self.build_runnable()
         variables = JupyternautVariables(
@@ -43,7 +43,7 @@ async def process_message(self, message: Message) -> None:
 
     def build_runnable(self) -> Any:
         # TODO: support model parameters. maybe we just add it to lm_provider_params in both 2.x and 3.x
-        llm = self.config.lm_provider(**self.config.lm_provider_params)
+        llm = self.config_manager.lm_provider(**self.config_manager.lm_provider_params)
         runnable = JUPYTERNAUT_PROMPT_TEMPLATE | llm | StrOutputParser()
 
         runnable = RunnableWithMessageHistory(

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

@@ -1,56 +1,92 @@
+from __future__ import annotations
+
 import asyncio
+import os
 from logging import Logger
 from time import time_ns
-from typing import TYPE_CHECKING, ClassVar, Optional
+from typing import TYPE_CHECKING, ClassVar
 
 from importlib_metadata import entry_points
 from jupyterlab_chat.models import Message
 from jupyterlab_chat.ychat import YChat
+from traitlets.config import LoggingConfigurable
 
 from ..config_manager import ConfigManager
 from .base_persona import BasePersona
 
 if TYPE_CHECKING:
     from asyncio import AbstractEventLoop
 
+    from jupyter_server_fileid.manager import (  # type: ignore[import-untyped]
+        BaseFileIdManager,
+    )
+
 # EPG := entry point group
 EPG_NAME = "jupyter_ai.personas"
 
 
-class PersonaManager:
+class PersonaManager(LoggingConfigurable):
     """
     Class that manages all personas for a single chat.
     """
 
     # instance attrs
     ychat: YChat
     config_manager: ConfigManager
-    event_loop: "AbstractEventLoop"
-    log: Logger
+    fileid_manager: BaseFileIdManager
+    root_dir: str
+    event_loop: AbstractEventLoop
+
+    log: Logger  # type: ignore
+    """
+    The `logging.Logger` instance used by this class. This is automatically set
+    by the `LoggingConfigurable` parent class; this declaration only hints the
+    type for type checkers.
+    """
+
     _personas: dict[str, BasePersona]
+    file_id: str
 
     # class attrs
-    _persona_classes: ClassVar[Optional[list[type[BasePersona]]]] = None
+    _persona_classes: ClassVar[list[type[BasePersona]] | None] = None
 
     def __init__(
         self,
+        *args,
+        room_id: str,
         ychat: YChat,
         config_manager: ConfigManager,
-        event_loop: "AbstractEventLoop",
-        log: Logger,
+        fileid_manager: BaseFileIdManager,
+        root_dir: str,
+        event_loop: AbstractEventLoop,
         message_interrupted: dict[str, asyncio.Event],
+        **kwargs,
     ):
+        # Forward other arguments to parent class
+        super().__init__(*args, **kwargs)
+
+        # Bind instance attributes
+        self.room_id = room_id
         self.ychat = ychat
         self.config_manager = config_manager
+        self.fileid_manager = fileid_manager
+        self.root_dir = root_dir
         self.event_loop = event_loop
-        self.log = log
         self.message_interrupted = message_interrupted
 
+        # Store file ID
+        self.file_id = room_id.split(":")[2]
+
+        # Load persona classes from entry points.
+        # This is stored in a class attribute (global to all instances) because
+        # the entry points are immutable after the server starts, so they only
+        # need to be loaded once.
         if not isinstance(PersonaManager._persona_classes, list):
             self._init_persona_classes()
             assert isinstance(PersonaManager._persona_classes, list)
 
         self._personas = self._init_personas()
+        self.log.error(self.get_chat_dir())
 
     def _init_persona_classes(self) -> None:
         """
@@ -124,10 +160,9 @@ def _init_personas(self) -> dict[str, BasePersona]:
         for Persona in persona_classes:
             try:
                 persona = Persona(
+                    parent=self,
                     ychat=self.ychat,
-                    manager=self,
-                    config=self.config_manager,
-                    log=self.log,
+                    config_manager=self.config_manager,
                     message_interrupted=self.message_interrupted,
                 )
             except Exception:
@@ -200,3 +235,28 @@ def route_message(self, new_message: Message):
         )
         for persona in mentioned_personas:
             self.event_loop.create_task(persona.process_message(new_message))
+
+    def get_chat_path(self, relative: bool = False) -> str:
+        """
+        Returns the absolute path of the chat file assigned to this
+        `PersonaManager`.
+
+        To get a path relative to the `ContentsManager` root directory, call
+        this method with `relative=True`.
+        """
+        relpath = self.fileid_manager.get_path(self.file_id)
+        if not relpath:
+            raise Exception(f"Unable to locate chat with file ID: '{self.file_id}'.")
+        if relative:
+            return relpath
+
+        abspath = os.path.join(self.root_dir, relpath)
+        return abspath
+
+    def get_chat_dir(self) -> str:
+        """
+        Returns the absolute path of the parent directory of the chat file
+        assigned to this `PersonaManager`.
+        """
+        abspath = self.get_chat_path(absolute=True)
+        return os.path.dirname(abspath)