From c464c1732d79ccca36ec0bd7558c184e44ddc233 Mon Sep 17 00:00:00 2001 From: Fernando Perez Date: Mon, 23 Jun 2025 15:35:01 -0700 Subject: [PATCH 01/13] First cut at the PersonaLoader imnplementation, no tests yet. --- .../jupyter_ai/personas/persona_manager.py | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index 05b63d4ff..d97c0d126 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -1,8 +1,12 @@ from __future__ import annotations import asyncio +import importlib.util +import inspect import os +from glob import glob from logging import Logger +from pathlib import Path from time import time_ns from typing import TYPE_CHECKING, Any, ClassVar @@ -140,6 +144,10 @@ def _init_persona_classes(self) -> None: "ERROR: Jupyter AI has no AI personas available. " + "Please verify your server configuration and open a new issue on our GitHub repo if this warning persists." ) + + # TODO: check whether in the end we need a full class for this, or if we can just use a simple function that returns a list of persona classes from the local filesystem. + persona_classes.extend(LocalPersonaLoader(self.root_dir).load_persona_classes()) + PersonaManager._persona_classes = persona_classes def _init_personas(self) -> dict[str, BasePersona]: @@ -287,3 +295,107 @@ def get_mcp_config(self) -> dict[str, Any]: return {} else: return self._mcp_config_loader.get_config(jdir) + + +class LocalPersonaLoader(LoggingConfigurable): + """ + Load _persona class declarations_ from the local filesystem. + + Those class declarations are then used to instantiate personas by the `PersonaManager`. + + TODO: wire to local .jupyter directory system, for now just use the current + directory. + """ + + def __init__( + self, + *args, + root_dir: str | None = None, + **kwargs, + ): + # Forward other arguments to parent class + super().__init__(*args, **kwargs) + + # Set default root directory to current working directory if not provided + self.root_dir = root_dir or os.getcwd() + + self.log.info(f"LocalPersonaLoader initialized with root directory: {self.root_dir}") + + def load_persona_classes(self) -> list[type[BasePersona]]: + """ + Loads persona classes from Python files in the local filesystem. + + Scans the root_dir for .py files, dynamically imports them, and extracts + any class declarations that are subclasses of BasePersona. + """ + persona_classes: list[type[BasePersona]] = [] + + # Check if root directory exists + if not os.path.exists(self.root_dir): + self.log.info(f"Root directory does not exist: {self.root_dir}") + return persona_classes + + # Find all .py files in the root directory + py_files = glob(os.path.join(self.root_dir, "*.py")) + + if not py_files: + self.log.info(f"No Python files found in directory: {self.root_dir}") + return persona_classes + + self.log.info(f"Found {len(py_files)} Python files in {self.root_dir}") + self.log.info("PENDING: Loading persona classes from local Python files...") + start_time_ns = time_ns() + + for py_file in py_files: + try: + # Get module name from file path + module_name = Path(py_file).stem + + # Skip if module name starts with underscore (private modules) + if module_name.startswith('_'): + continue + + # Create module spec and load the module + spec = importlib.util.spec_from_file_location(module_name, py_file) + if spec is None or spec.loader is None: + self.log.warning(f" - Unable to create module spec for {py_file}") + continue + + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + # Find all classes in the module that are BasePersona subclasses + module_persona_classes = [] + for name, obj in inspect.getmembers(module, inspect.isclass): + # Check if it's a subclass of BasePersona but not BasePersona itself + if (issubclass(obj, BasePersona) and + obj is not BasePersona and + obj.__module__ == module_name): + module_persona_classes.append(obj) + + if module_persona_classes: + persona_classes.extend(module_persona_classes) + class_names = [cls.__name__ for cls in module_persona_classes] + self.log.info( + f" - Loaded {len(module_persona_classes)} persona class(es) from '{py_file}': {class_names}" + ) + else: + self.log.debug(f" - No persona classes found in '{py_file}'") + + except Exception: + # On exception, log an error and continue + # This mirrors the error handling pattern from entry point loading + self.log.exception( + f" - Unable to load persona classes from '{py_file}' due to an exception printed below." + ) + continue + + if len(persona_classes) > 0: + elapsed_time_ms = (time_ns() - start_time_ns) // 1_000_000 + self.log.info( + f"SUCCESS: Loaded {len(persona_classes)} persona classes from local filesystem. Time elapsed: {elapsed_time_ms}ms." + ) + else: + self.log.info("No persona classes found in local filesystem.") + + return persona_classes From 241a53023755c24afe950a53cfc04c95e94a21fd Mon Sep 17 00:00:00 2001 From: Fernando Perez Date: Mon, 23 Jun 2025 15:38:44 -0700 Subject: [PATCH 02/13] Scope import logic more tightly for persona files. Only match (case insensitively) Python files with `persona` in the name. --- .../jupyter-ai/jupyter_ai/personas/persona_manager.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index d97c0d126..46996b607 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -335,14 +335,15 @@ def load_persona_classes(self) -> list[type[BasePersona]]: self.log.info(f"Root directory does not exist: {self.root_dir}") return persona_classes - # Find all .py files in the root directory - py_files = glob(os.path.join(self.root_dir, "*.py")) + # Find all .py files in the root directory that contain "persona" in the name + all_py_files = glob(os.path.join(self.root_dir, "*.py")) + py_files = [f for f in all_py_files if "persona" in Path(f).stem.lower()] if not py_files: - self.log.info(f"No Python files found in directory: {self.root_dir}") + self.log.info(f"No Python files with 'persona' in the name found in directory: {self.root_dir}") return persona_classes - self.log.info(f"Found {len(py_files)} Python files in {self.root_dir}") + self.log.info(f"Found {len(py_files)} Python files with 'persona' in the name in {self.root_dir}") self.log.info("PENDING: Loading persona classes from local Python files...") start_time_ns = time_ns() From e3edf80de5655e5dbe9b075144694a08da513816 Mon Sep 17 00:00:00 2001 From: Fernando Perez Date: Mon, 23 Jun 2025 18:14:31 -0700 Subject: [PATCH 03/13] Add first pass of unit tests for the local persona loader. --- .../jupyter_ai/tests/test_personas.py | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) create mode 100644 packages/jupyter-ai/jupyter_ai/tests/test_personas.py diff --git a/packages/jupyter-ai/jupyter_ai/tests/test_personas.py b/packages/jupyter-ai/jupyter_ai/tests/test_personas.py new file mode 100644 index 000000000..3c60f9655 --- /dev/null +++ b/packages/jupyter-ai/jupyter_ai/tests/test_personas.py @@ -0,0 +1,77 @@ +""" +Test the local persona manager. +""" + +import pytest +import tempfile +from pathlib import Path + +from jupyter_ai.personas.base_persona import BasePersona, PersonaDefaults +from jupyter_ai.personas.persona_manager import LocalPersonaLoader + + +@pytest.fixture +def tmp_persona_dir(): + """Create a temporary directory for testing LocalPersonaLoader with guaranteed cleanup.""" + with tempfile.TemporaryDirectory() as temp_dir: + yield Path(temp_dir) + + +class TestLocalPersonaLoader: + """Test cases for LocalPersonaLoader class.""" + + def test_empty_directory_returns_empty_list(self, tmp_persona_dir): + """Test that an empty directory returns an empty list of persona classes.""" + loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir)) + result = loader.load_persona_classes() + assert result == [] + + def test_non_persona_file_returns_empty_list(self, tmp_persona_dir): + """Test that a Python file without persona classes returns an empty list.""" + # Create a file that doesn't contain "persona" in the name + non_persona_file = tmp_persona_dir / "no_personas.py" + non_persona_file.write_text("pass") + + loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir)) + result = loader.load_persona_classes() + assert result == [] + + def test_simple_persona_file_returns_persona_class(self, tmp_persona_dir): + """Test that a file with a BasePersona subclass returns that class.""" + # Create a simple persona file + persona_file = tmp_persona_dir / "simple_personas.py" + persona_content = ''' +from jupyter_ai.personas.base_persona import BasePersona + +class TestPersona(BasePersona): + id = "test_persona" + name = "Test Persona" + description = "A simple test persona" + + def process_message(self, message): + pass +''' + persona_file.write_text(persona_content) + + loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir)) + result = loader.load_persona_classes() + + assert len(result) == 1 + assert result[0].__name__ == "TestPersona" + assert issubclass(result[0], BasePersona) + + def test_bad_persona_file_logs_exception_returns_empty_list(self, tmp_persona_dir, caplog): + """Test that a file with syntax errors logs an exception and returns empty list.""" + # Create a file with invalid Python code + bad_persona_file = tmp_persona_dir / "bad_persona.py" + bad_persona_file.write_text("1/0") + + loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir)) + result = loader.load_persona_classes() + + assert result == [] + # Check that an error was logged + assert any("Unable to load persona classes from" in record.message + for record in caplog.records if record.levelname in ["ERROR", "EXCEPTION"]) + + From 8ca434e04eedebd8b0df4eca80175c121aed39e5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 24 Jun 2025 01:18:01 +0000 Subject: [PATCH 04/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../jupyter_ai/personas/persona_manager.py | 54 +++++++++++-------- .../jupyter_ai/tests/test_personas.py | 41 +++++++------- 2 files changed, 53 insertions(+), 42 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index 46996b607..fc3846353 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -315,65 +315,73 @@ def __init__( ): # Forward other arguments to parent class super().__init__(*args, **kwargs) - + # Set default root directory to current working directory if not provided self.root_dir = root_dir or os.getcwd() - - self.log.info(f"LocalPersonaLoader initialized with root directory: {self.root_dir}") + + self.log.info( + f"LocalPersonaLoader initialized with root directory: {self.root_dir}" + ) def load_persona_classes(self) -> list[type[BasePersona]]: """ Loads persona classes from Python files in the local filesystem. - + Scans the root_dir for .py files, dynamically imports them, and extracts any class declarations that are subclasses of BasePersona. """ persona_classes: list[type[BasePersona]] = [] - + # Check if root directory exists if not os.path.exists(self.root_dir): self.log.info(f"Root directory does not exist: {self.root_dir}") return persona_classes - + # Find all .py files in the root directory that contain "persona" in the name all_py_files = glob(os.path.join(self.root_dir, "*.py")) py_files = [f for f in all_py_files if "persona" in Path(f).stem.lower()] - + if not py_files: - self.log.info(f"No Python files with 'persona' in the name found in directory: {self.root_dir}") + self.log.info( + f"No Python files with 'persona' in the name found in directory: {self.root_dir}" + ) return persona_classes - - self.log.info(f"Found {len(py_files)} Python files with 'persona' in the name in {self.root_dir}") + + self.log.info( + f"Found {len(py_files)} Python files with 'persona' in the name in {self.root_dir}" + ) self.log.info("PENDING: Loading persona classes from local Python files...") start_time_ns = time_ns() - + for py_file in py_files: try: # Get module name from file path module_name = Path(py_file).stem - + # Skip if module name starts with underscore (private modules) - if module_name.startswith('_'): + if module_name.startswith("_"): continue - + # Create module spec and load the module spec = importlib.util.spec_from_file_location(module_name, py_file) if spec is None or spec.loader is None: self.log.warning(f" - Unable to create module spec for {py_file}") continue - + module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) - + # Find all classes in the module that are BasePersona subclasses module_persona_classes = [] for name, obj in inspect.getmembers(module, inspect.isclass): # Check if it's a subclass of BasePersona but not BasePersona itself - if (issubclass(obj, BasePersona) and - obj is not BasePersona and - obj.__module__ == module_name): + if ( + issubclass(obj, BasePersona) + and obj is not BasePersona + and obj.__module__ == module_name + ): module_persona_classes.append(obj) - + if module_persona_classes: persona_classes.extend(module_persona_classes) class_names = [cls.__name__ for cls in module_persona_classes] @@ -382,7 +390,7 @@ def load_persona_classes(self) -> list[type[BasePersona]]: ) else: self.log.debug(f" - No persona classes found in '{py_file}'") - + except Exception: # On exception, log an error and continue # This mirrors the error handling pattern from entry point loading @@ -390,7 +398,7 @@ def load_persona_classes(self) -> list[type[BasePersona]]: f" - Unable to load persona classes from '{py_file}' due to an exception printed below." ) continue - + if len(persona_classes) > 0: elapsed_time_ms = (time_ns() - start_time_ns) // 1_000_000 self.log.info( @@ -398,5 +406,5 @@ def load_persona_classes(self) -> list[type[BasePersona]]: ) else: self.log.info("No persona classes found in local filesystem.") - + return persona_classes diff --git a/packages/jupyter-ai/jupyter_ai/tests/test_personas.py b/packages/jupyter-ai/jupyter_ai/tests/test_personas.py index 3c60f9655..c210ae178 100644 --- a/packages/jupyter-ai/jupyter_ai/tests/test_personas.py +++ b/packages/jupyter-ai/jupyter_ai/tests/test_personas.py @@ -2,10 +2,10 @@ Test the local persona manager. """ -import pytest import tempfile from pathlib import Path +import pytest from jupyter_ai.personas.base_persona import BasePersona, PersonaDefaults from jupyter_ai.personas.persona_manager import LocalPersonaLoader @@ -19,59 +19,62 @@ def tmp_persona_dir(): class TestLocalPersonaLoader: """Test cases for LocalPersonaLoader class.""" - + def test_empty_directory_returns_empty_list(self, tmp_persona_dir): """Test that an empty directory returns an empty list of persona classes.""" loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir)) result = loader.load_persona_classes() assert result == [] - + def test_non_persona_file_returns_empty_list(self, tmp_persona_dir): """Test that a Python file without persona classes returns an empty list.""" # Create a file that doesn't contain "persona" in the name non_persona_file = tmp_persona_dir / "no_personas.py" non_persona_file.write_text("pass") - + loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir)) result = loader.load_persona_classes() assert result == [] - + def test_simple_persona_file_returns_persona_class(self, tmp_persona_dir): """Test that a file with a BasePersona subclass returns that class.""" # Create a simple persona file persona_file = tmp_persona_dir / "simple_personas.py" - persona_content = ''' + persona_content = """ from jupyter_ai.personas.base_persona import BasePersona class TestPersona(BasePersona): id = "test_persona" name = "Test Persona" description = "A simple test persona" - + def process_message(self, message): pass -''' +""" persona_file.write_text(persona_content) - + loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir)) result = loader.load_persona_classes() - + assert len(result) == 1 assert result[0].__name__ == "TestPersona" assert issubclass(result[0], BasePersona) - - def test_bad_persona_file_logs_exception_returns_empty_list(self, tmp_persona_dir, caplog): + + def test_bad_persona_file_logs_exception_returns_empty_list( + self, tmp_persona_dir, caplog + ): """Test that a file with syntax errors logs an exception and returns empty list.""" # Create a file with invalid Python code - bad_persona_file = tmp_persona_dir / "bad_persona.py" + bad_persona_file = tmp_persona_dir / "bad_persona.py" bad_persona_file.write_text("1/0") - + loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir)) result = loader.load_persona_classes() - + assert result == [] # Check that an error was logged - assert any("Unable to load persona classes from" in record.message - for record in caplog.records if record.levelname in ["ERROR", "EXCEPTION"]) - - + assert any( + "Unable to load persona classes from" in record.message + for record in caplog.records + if record.levelname in ["ERROR", "EXCEPTION"] + ) From af2b6a5e98d94cc5ca88e0c7c997499c20a2dae9 Mon Sep 17 00:00:00 2001 From: Fernando Perez Date: Thu, 26 Jun 2025 10:09:35 -0700 Subject: [PATCH 05/13] Refactor into standalone function --- .../jupyter_ai/personas/persona_manager.py | 152 ++++++------------ .../jupyter_ai/tests/test_personas.py | 30 ++-- 2 files changed, 61 insertions(+), 121 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index fc3846353..70785551b 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -145,8 +145,8 @@ def _init_persona_classes(self) -> None: + "Please verify your server configuration and open a new issue on our GitHub repo if this warning persists." ) - # TODO: check whether in the end we need a full class for this, or if we can just use a simple function that returns a list of persona classes from the local filesystem. - persona_classes.extend(LocalPersonaLoader(self.root_dir).load_persona_classes()) + # Load persona classes from local filesystem + persona_classes.extend(load_persona_classes_from_directory(self.root_dir)) PersonaManager._persona_classes = persona_classes @@ -297,114 +297,66 @@ def get_mcp_config(self) -> dict[str, Any]: return self._mcp_config_loader.get_config(jdir) -class LocalPersonaLoader(LoggingConfigurable): +def load_persona_classes_from_directory(root_dir: str) -> list[type[BasePersona]]: """ - Load _persona class declarations_ from the local filesystem. + Load _persona class declarations_ from Python files in the local filesystem. - Those class declarations are then used to instantiate personas by the `PersonaManager`. + Those class declarations are then used to instantiate personas by the + `PersonaManager`. - TODO: wire to local .jupyter directory system, for now just use the current - directory. - """ - - def __init__( - self, - *args, - root_dir: str | None = None, - **kwargs, - ): - # Forward other arguments to parent class - super().__init__(*args, **kwargs) + Scans the root_dir for .py files containing "persona" in their name, + dynamically imports them, and extracts any class declarations that are + subclasses of BasePersona. - # Set default root directory to current working directory if not provided - self.root_dir = root_dir or os.getcwd() + Args: + root_dir: Directory to scan for persona Python files - self.log.info( - f"LocalPersonaLoader initialized with root directory: {self.root_dir}" - ) - - def load_persona_classes(self) -> list[type[BasePersona]]: - """ - Loads persona classes from Python files in the local filesystem. - - Scans the root_dir for .py files, dynamically imports them, and extracts - any class declarations that are subclasses of BasePersona. - """ - persona_classes: list[type[BasePersona]] = [] - - # Check if root directory exists - if not os.path.exists(self.root_dir): - self.log.info(f"Root directory does not exist: {self.root_dir}") - return persona_classes + Returns: + List of BasePersona subclasses found in the directory + """ + persona_classes: list[type[BasePersona]] = [] - # Find all .py files in the root directory that contain "persona" in the name - all_py_files = glob(os.path.join(self.root_dir, "*.py")) - py_files = [f for f in all_py_files if "persona" in Path(f).stem.lower()] + # Check if root directory exists + if not os.path.exists(root_dir): + return persona_classes - if not py_files: - self.log.info( - f"No Python files with 'persona' in the name found in directory: {self.root_dir}" - ) - return persona_classes + # Find all .py files in the root directory that contain "persona" in the name + all_py_files = glob(os.path.join(root_dir, "*.py")) + py_files = [f for f in all_py_files if "persona" in Path(f).stem.lower()] - self.log.info( - f"Found {len(py_files)} Python files with 'persona' in the name in {self.root_dir}" - ) - self.log.info("PENDING: Loading persona classes from local Python files...") - start_time_ns = time_ns() + if not py_files: + return persona_classes - for py_file in py_files: - try: - # Get module name from file path - module_name = Path(py_file).stem - - # Skip if module name starts with underscore (private modules) - if module_name.startswith("_"): - continue - - # Create module spec and load the module - spec = importlib.util.spec_from_file_location(module_name, py_file) - if spec is None or spec.loader is None: - self.log.warning(f" - Unable to create module spec for {py_file}") - continue - - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - - # Find all classes in the module that are BasePersona subclasses - module_persona_classes = [] - for name, obj in inspect.getmembers(module, inspect.isclass): - # Check if it's a subclass of BasePersona but not BasePersona itself - if ( - issubclass(obj, BasePersona) - and obj is not BasePersona - and obj.__module__ == module_name - ): - module_persona_classes.append(obj) - - if module_persona_classes: - persona_classes.extend(module_persona_classes) - class_names = [cls.__name__ for cls in module_persona_classes] - self.log.info( - f" - Loaded {len(module_persona_classes)} persona class(es) from '{py_file}': {class_names}" - ) - else: - self.log.debug(f" - No persona classes found in '{py_file}'") + for py_file in py_files: + try: + # Get module name from file path + module_name = Path(py_file).stem - except Exception: - # On exception, log an error and continue - # This mirrors the error handling pattern from entry point loading - self.log.exception( - f" - Unable to load persona classes from '{py_file}' due to an exception printed below." - ) + # Skip if module name starts with underscore (private modules) + if module_name.startswith("_"): continue - if len(persona_classes) > 0: - elapsed_time_ms = (time_ns() - start_time_ns) // 1_000_000 - self.log.info( - f"SUCCESS: Loaded {len(persona_classes)} persona classes from local filesystem. Time elapsed: {elapsed_time_ms}ms." - ) - else: - self.log.info("No persona classes found in local filesystem.") + # Create module spec and load the module + spec = importlib.util.spec_from_file_location(module_name, py_file) + if spec is None or spec.loader is None: + continue - return persona_classes + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + # Find all classes in the module that are BasePersona subclasses + for name, obj in inspect.getmembers(module, inspect.isclass): + # Check if it's a subclass of BasePersona but not BasePersona itself + if ( + issubclass(obj, BasePersona) + and obj is not BasePersona + and obj.__module__ == module_name + ): + persona_classes.append(obj) + + except Exception: + # On exception, continue to next file + # This mirrors the error handling pattern from entry point loading + continue + + return persona_classes diff --git a/packages/jupyter-ai/jupyter_ai/tests/test_personas.py b/packages/jupyter-ai/jupyter_ai/tests/test_personas.py index c210ae178..ae18c9232 100644 --- a/packages/jupyter-ai/jupyter_ai/tests/test_personas.py +++ b/packages/jupyter-ai/jupyter_ai/tests/test_personas.py @@ -7,7 +7,7 @@ import pytest from jupyter_ai.personas.base_persona import BasePersona, PersonaDefaults -from jupyter_ai.personas.persona_manager import LocalPersonaLoader +from jupyter_ai.personas.persona_manager import load_persona_classes_from_directory @pytest.fixture @@ -17,13 +17,12 @@ def tmp_persona_dir(): yield Path(temp_dir) -class TestLocalPersonaLoader: - """Test cases for LocalPersonaLoader class.""" +class TestLoadPersonaClassesFromDirectory: + """Test cases for load_persona_classes_from_directory function.""" def test_empty_directory_returns_empty_list(self, tmp_persona_dir): """Test that an empty directory returns an empty list of persona classes.""" - loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir)) - result = loader.load_persona_classes() + result = load_persona_classes_from_directory(str(tmp_persona_dir)) assert result == [] def test_non_persona_file_returns_empty_list(self, tmp_persona_dir): @@ -32,8 +31,7 @@ def test_non_persona_file_returns_empty_list(self, tmp_persona_dir): non_persona_file = tmp_persona_dir / "no_personas.py" non_persona_file.write_text("pass") - loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir)) - result = loader.load_persona_classes() + result = load_persona_classes_from_directory(str(tmp_persona_dir)) assert result == [] def test_simple_persona_file_returns_persona_class(self, tmp_persona_dir): @@ -53,28 +51,18 @@ def process_message(self, message): """ persona_file.write_text(persona_content) - loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir)) - result = loader.load_persona_classes() + result = load_persona_classes_from_directory(str(tmp_persona_dir)) assert len(result) == 1 assert result[0].__name__ == "TestPersona" assert issubclass(result[0], BasePersona) - def test_bad_persona_file_logs_exception_returns_empty_list( - self, tmp_persona_dir, caplog - ): - """Test that a file with syntax errors logs an exception and returns empty list.""" + def test_bad_persona_file_returns_empty_list(self, tmp_persona_dir): + """Test that a file with syntax errors returns empty list.""" # Create a file with invalid Python code bad_persona_file = tmp_persona_dir / "bad_persona.py" bad_persona_file.write_text("1/0") - loader = LocalPersonaLoader(root_dir=str(tmp_persona_dir)) - result = loader.load_persona_classes() + result = load_persona_classes_from_directory(str(tmp_persona_dir)) assert result == [] - # Check that an error was logged - assert any( - "Unable to load persona classes from" in record.message - for record in caplog.records - if record.levelname in ["ERROR", "EXCEPTION"] - ) From fe69bd8ae32bfdb4b63ff969f4d065ff132ac2cf Mon Sep 17 00:00:00 2001 From: Fernando Perez Date: Thu, 26 Jun 2025 19:58:48 -0700 Subject: [PATCH 06/13] Make persona_classes an instance attr so it actually reloads per chat file. Also, rename local loader to a shorter name. --- .../jupyter_ai/personas/persona_manager.py | 44 +++++++++++-------- .../jupyter_ai/tests/test_personas.py | 12 ++--- 2 files changed, 31 insertions(+), 25 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index 70785551b..0eed64c5a 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -51,11 +51,10 @@ class PersonaManager(LoggingConfigurable): type for type checkers. """ + _persona_classes: ClassVar[list[type[BasePersona]] | None] = None _personas: dict[str, BasePersona] file_id: str - # class attrs - _persona_classes: ClassVar[list[type[BasePersona]] | None] = None def __init__( self, @@ -91,9 +90,9 @@ def __init__( # 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): + if not isinstance(self._persona_classes, list): self._init_persona_classes() - assert isinstance(PersonaManager._persona_classes, list) + assert isinstance(self._persona_classes, list) self._personas = self._init_personas() @@ -102,12 +101,18 @@ def _init_persona_classes(self) -> None: Initializes the list of persona *classes* by retrieving the `jupyter-ai.personas` entry points group. - This list is cached in the `PersonaManager._persona_classes` class - attribute, i.e. this method should only run once in the extension + # TODO: fix this part of docs now that we have it as an instance attr. + This list is cached in the `self._persona_classes` instance + attribute, .e. this method should only run once in the extension lifecycle. """ - if PersonaManager._persona_classes: - return + # Loading is in two parts: + # 1. Load persona classes from package entry points. + # 2. Load persona classes from local filesystem. + # + # This allows for lightweight development of new personas by the user in + # their local filesystem. The first part is done here, and the second + # part is done in `_init_personas()`. persona_eps = entry_points().select(group=EPG_NAME) self.log.info(f"Found {len(persona_eps)} entry points under '{EPG_NAME}'.") @@ -146,9 +151,9 @@ def _init_persona_classes(self) -> None: ) # Load persona classes from local filesystem - persona_classes.extend(load_persona_classes_from_directory(self.root_dir)) + persona_classes.extend(load_from_dir(self.get_dotjupyter_dir())) - PersonaManager._persona_classes = persona_classes + self._persona_classes = persona_classes def _init_personas(self) -> dict[str, BasePersona]: """ @@ -156,7 +161,7 @@ def _init_personas(self) -> dict[str, BasePersona]: to the constructor. """ # Ensure that persona classes were initialized first - persona_classes = PersonaManager._persona_classes + persona_classes = self._persona_classes assert isinstance(persona_classes, list) # If no persona classes are available, log a warning and return @@ -297,22 +302,23 @@ def get_mcp_config(self) -> dict[str, Any]: return self._mcp_config_loader.get_config(jdir) -def load_persona_classes_from_directory(root_dir: str) -> list[type[BasePersona]]: +def load_from_dir(root_dir: str) -> list[type[BasePersona]]: """ Load _persona class declarations_ from Python files in the local filesystem. Those class declarations are then used to instantiate personas by the `PersonaManager`. - Scans the root_dir for .py files containing "persona" in their name, + Scans the root_dir for .py files containing `persona` in their name that do + _not_ start with a single `_` (i.e. private modules are skipped). Then, it dynamically imports them, and extracts any class declarations that are - subclasses of BasePersona. + subclasses of `BasePersona`. Args: - root_dir: Directory to scan for persona Python files + root_dir: Directory to scan for persona Python files. Returns: - List of BasePersona subclasses found in the directory + List of `BasePersona` subclasses found in the directory. """ persona_classes: list[type[BasePersona]] = [] @@ -321,12 +327,12 @@ def load_persona_classes_from_directory(root_dir: str) -> list[type[BasePersona] return persona_classes # Find all .py files in the root directory that contain "persona" in the name + # TODO: protect with a try/except block for exceptions raised by glob all_py_files = glob(os.path.join(root_dir, "*.py")) py_files = [f for f in all_py_files if "persona" in Path(f).stem.lower()] - if not py_files: - return persona_classes - + # For each .py file, dynamically import the module and extract all + # BasePersona subclasses. for py_file in py_files: try: # Get module name from file path diff --git a/packages/jupyter-ai/jupyter_ai/tests/test_personas.py b/packages/jupyter-ai/jupyter_ai/tests/test_personas.py index ae18c9232..8f3aa290f 100644 --- a/packages/jupyter-ai/jupyter_ai/tests/test_personas.py +++ b/packages/jupyter-ai/jupyter_ai/tests/test_personas.py @@ -7,7 +7,7 @@ import pytest from jupyter_ai.personas.base_persona import BasePersona, PersonaDefaults -from jupyter_ai.personas.persona_manager import load_persona_classes_from_directory +from jupyter_ai.personas.persona_manager import load_from_dir @pytest.fixture @@ -18,11 +18,11 @@ def tmp_persona_dir(): class TestLoadPersonaClassesFromDirectory: - """Test cases for load_persona_classes_from_directory function.""" + """Test cases for load_from_dir function.""" def test_empty_directory_returns_empty_list(self, tmp_persona_dir): """Test that an empty directory returns an empty list of persona classes.""" - result = load_persona_classes_from_directory(str(tmp_persona_dir)) + result = load_from_dir(str(tmp_persona_dir)) assert result == [] def test_non_persona_file_returns_empty_list(self, tmp_persona_dir): @@ -31,7 +31,7 @@ def test_non_persona_file_returns_empty_list(self, tmp_persona_dir): non_persona_file = tmp_persona_dir / "no_personas.py" non_persona_file.write_text("pass") - result = load_persona_classes_from_directory(str(tmp_persona_dir)) + result = load_from_dir(str(tmp_persona_dir)) assert result == [] def test_simple_persona_file_returns_persona_class(self, tmp_persona_dir): @@ -51,7 +51,7 @@ def process_message(self, message): """ persona_file.write_text(persona_content) - result = load_persona_classes_from_directory(str(tmp_persona_dir)) + result = load_from_dir(str(tmp_persona_dir)) assert len(result) == 1 assert result[0].__name__ == "TestPersona" @@ -63,6 +63,6 @@ def test_bad_persona_file_returns_empty_list(self, tmp_persona_dir): bad_persona_file = tmp_persona_dir / "bad_persona.py" bad_persona_file.write_text("1/0") - result = load_persona_classes_from_directory(str(tmp_persona_dir)) + result = load_from_dir(str(tmp_persona_dir)) assert result == [] From ff1ad2fdcd566f8f373768c7554a63e836623fd1 Mon Sep 17 00:00:00 2001 From: Fernando Perez Date: Thu, 26 Jun 2025 20:39:06 -0700 Subject: [PATCH 07/13] Improve logging so we can debug persona loading issues --- .../jupyter_ai/personas/persona_manager.py | 32 +++++++++----- .../jupyter_ai/tests/test_personas.py | 43 +++++++++++-------- 2 files changed, 46 insertions(+), 29 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index 0eed64c5a..467ae5809 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -151,7 +151,7 @@ def _init_persona_classes(self) -> None: ) # Load persona classes from local filesystem - persona_classes.extend(load_from_dir(self.get_dotjupyter_dir())) + persona_classes.extend(load_from_dir(self.get_dotjupyter_dir(), self.log)) self._persona_classes = persona_classes @@ -302,7 +302,7 @@ def get_mcp_config(self) -> dict[str, Any]: return self._mcp_config_loader.get_config(jdir) -def load_from_dir(root_dir: str) -> list[type[BasePersona]]: +def load_from_dir(root_dir: str, log: Logger) -> list[type[BasePersona]]: """ Load _persona class declarations_ from Python files in the local filesystem. @@ -314,11 +314,12 @@ def load_from_dir(root_dir: str) -> list[type[BasePersona]]: dynamically imports them, and extracts any class declarations that are subclasses of `BasePersona`. - Args: - root_dir: Directory to scan for persona Python files. + Args: + root_dir: Directory to scan for persona Python files. + log: Logger instance for logging messages. - Returns: - List of `BasePersona` subclasses found in the directory. + Returns: + List of `BasePersona` subclasses found in the directory. """ persona_classes: list[type[BasePersona]] = [] @@ -327,10 +328,17 @@ def load_from_dir(root_dir: str) -> list[type[BasePersona]]: return persona_classes # Find all .py files in the root directory that contain "persona" in the name - # TODO: protect with a try/except block for exceptions raised by glob - all_py_files = glob(os.path.join(root_dir, "*.py")) - py_files = [f for f in all_py_files if "persona" in Path(f).stem.lower()] + try: + all_py_files = glob(os.path.join(root_dir, "*.py")) + py_files = [f for f in all_py_files if "persona" in Path(f).stem.lower()] + except Exception as e: + # On exception with glob operation, return empty list + log.debug(f"{type(e).__name__} occurred while searching for Python files in {root_dir}") + return persona_classes + if py_files: + log.info(f"Loading persona files from {root_dir}: {[Path(f).name for f in py_files]}") + # For each .py file, dynamically import the module and extract all # BasePersona subclasses. for py_file in py_files: @@ -360,9 +368,11 @@ def load_from_dir(root_dir: str) -> list[type[BasePersona]]: ): persona_classes.append(obj) - except Exception: - # On exception, continue to next file + except Exception as e: + # On exception, log error and continue to next file # This mirrors the error handling pattern from entry point loading + log.error(f"Unable to load persona classes from '{py_file}'") + log.error(f"Error was: {type(e).__name__}: {e} - Unable to load persona classes from '{py_file}'") continue return persona_classes diff --git a/packages/jupyter-ai/jupyter_ai/tests/test_personas.py b/packages/jupyter-ai/jupyter_ai/tests/test_personas.py index 8f3aa290f..17bea6532 100644 --- a/packages/jupyter-ai/jupyter_ai/tests/test_personas.py +++ b/packages/jupyter-ai/jupyter_ai/tests/test_personas.py @@ -4,6 +4,7 @@ import tempfile from pathlib import Path +from unittest.mock import Mock import pytest from jupyter_ai.personas.base_persona import BasePersona, PersonaDefaults @@ -17,24 +18,30 @@ def tmp_persona_dir(): yield Path(temp_dir) +@pytest.fixture +def mock_logger(): + """Create a mock logger for testing.""" + return Mock() + + class TestLoadPersonaClassesFromDirectory: """Test cases for load_from_dir function.""" - def test_empty_directory_returns_empty_list(self, tmp_persona_dir): + def test_empty_directory_returns_empty_list(self, tmp_persona_dir, mock_logger): """Test that an empty directory returns an empty list of persona classes.""" - result = load_from_dir(str(tmp_persona_dir)) + result = load_from_dir(str(tmp_persona_dir), mock_logger) assert result == [] - - def test_non_persona_file_returns_empty_list(self, tmp_persona_dir): + + def test_non_persona_file_returns_empty_list(self, tmp_persona_dir, mock_logger): """Test that a Python file without persona classes returns an empty list.""" # Create a file that doesn't contain "persona" in the name non_persona_file = tmp_persona_dir / "no_personas.py" non_persona_file.write_text("pass") - - result = load_from_dir(str(tmp_persona_dir)) + + result = load_from_dir(str(tmp_persona_dir), mock_logger) assert result == [] - - def test_simple_persona_file_returns_persona_class(self, tmp_persona_dir): + + def test_simple_persona_file_returns_persona_class(self, tmp_persona_dir, mock_logger): """Test that a file with a BasePersona subclass returns that class.""" # Create a simple persona file persona_file = tmp_persona_dir / "simple_personas.py" @@ -45,24 +52,24 @@ class TestPersona(BasePersona): id = "test_persona" name = "Test Persona" description = "A simple test persona" - + def process_message(self, message): pass """ persona_file.write_text(persona_content) - - result = load_from_dir(str(tmp_persona_dir)) - + + result = load_from_dir(str(tmp_persona_dir), mock_logger) + assert len(result) == 1 assert result[0].__name__ == "TestPersona" assert issubclass(result[0], BasePersona) - - def test_bad_persona_file_returns_empty_list(self, tmp_persona_dir): + + def test_bad_persona_file_returns_empty_list(self, tmp_persona_dir, mock_logger): """Test that a file with syntax errors returns empty list.""" # Create a file with invalid Python code - bad_persona_file = tmp_persona_dir / "bad_persona.py" + bad_persona_file = tmp_persona_dir / "bad_persona.py" bad_persona_file.write_text("1/0") - - result = load_from_dir(str(tmp_persona_dir)) - + + result = load_from_dir(str(tmp_persona_dir), mock_logger) + assert result == [] From 251728618e364aa0b3cde5b55621cc01b0421759 Mon Sep 17 00:00:00 2001 From: Fernando Perez Date: Thu, 26 Jun 2025 21:34:22 -0700 Subject: [PATCH 08/13] Ignore files with leading dot --- packages/jupyter-ai/jupyter_ai/personas/persona_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index 467ae5809..bf5c9d324 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -347,7 +347,7 @@ def load_from_dir(root_dir: str, log: Logger) -> list[type[BasePersona]]: module_name = Path(py_file).stem # Skip if module name starts with underscore (private modules) - if module_name.startswith("_"): + if module_name.startswith("_") or module_name.startswith("."): continue # Create module spec and load the module From 279096112d5b41ab0e7c1f2b7cd4d3ee7094091c Mon Sep 17 00:00:00 2001 From: Fernando Perez Date: Thu, 26 Jun 2025 23:02:42 -0700 Subject: [PATCH 09/13] Add a bit more logging for local persona loading --- packages/jupyter-ai/jupyter_ai/personas/persona_manager.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index bf5c9d324..9a31a87df 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -323,6 +323,7 @@ def load_from_dir(root_dir: str, log: Logger) -> list[type[BasePersona]]: """ persona_classes: list[type[BasePersona]] = [] + log.info(f"Loading persona files from {root_dir}") # Check if root directory exists if not os.path.exists(root_dir): return persona_classes @@ -366,6 +367,7 @@ def load_from_dir(root_dir: str, log: Logger) -> list[type[BasePersona]]: and obj is not BasePersona and obj.__module__ == module_name ): + log.info(f"Found persona class '{obj.__name__}' in '{py_file}'") persona_classes.append(obj) except Exception as e: @@ -376,3 +378,4 @@ def load_from_dir(root_dir: str, log: Logger) -> list[type[BasePersona]]: continue return persona_classes + From d0d97cf4a55892c0becc44ad9cb53ee1d32e7bfb Mon Sep 17 00:00:00 2001 From: Fernando Perez Date: Thu, 26 Jun 2025 23:10:06 -0700 Subject: [PATCH 10/13] Fix logic to skip . and _ files from the start instead of in the later loop, clarify logging --- .../jupyter_ai/personas/persona_manager.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index 9a31a87df..12645939c 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -323,7 +323,7 @@ def load_from_dir(root_dir: str, log: Logger) -> list[type[BasePersona]]: """ persona_classes: list[type[BasePersona]] = [] - log.info(f"Loading persona files from {root_dir}") + log.info(f"Searching for persona files in {root_dir}") # Check if root directory exists if not os.path.exists(root_dir): return persona_classes @@ -331,14 +331,19 @@ def load_from_dir(root_dir: str, log: Logger) -> list[type[BasePersona]]: # Find all .py files in the root directory that contain "persona" in the name try: all_py_files = glob(os.path.join(root_dir, "*.py")) - py_files = [f for f in all_py_files if "persona" in Path(f).stem.lower()] + py_files = [] + for f in all_py_files: + fname_lower = Path(f).stem.lower() + if "persona" in fname_lower and not (fname_lower.startswith("_") or fname_lower.startswith(".")): + py_files.append(f) + except Exception as e: # On exception with glob operation, return empty list - log.debug(f"{type(e).__name__} occurred while searching for Python files in {root_dir}") + log.error(f"{type(e).__name__} occurred while searching for Python files in {root_dir}") return persona_classes if py_files: - log.info(f"Loading persona files from {root_dir}: {[Path(f).name for f in py_files]}") + log.info(f"Found files from {root_dir}: {[Path(f).name for f in py_files]}") # For each .py file, dynamically import the module and extract all # BasePersona subclasses. @@ -347,10 +352,6 @@ def load_from_dir(root_dir: str, log: Logger) -> list[type[BasePersona]]: # Get module name from file path module_name = Path(py_file).stem - # Skip if module name starts with underscore (private modules) - if module_name.startswith("_") or module_name.startswith("."): - continue - # Create module spec and load the module spec = importlib.util.spec_from_file_location(module_name, py_file) if spec is None or spec.loader is None: From 1f9bcaff379799a552096176cf9cdd3104c29b09 Mon Sep 17 00:00:00 2001 From: Fernando Perez Date: Thu, 26 Jun 2025 23:15:40 -0700 Subject: [PATCH 11/13] Use log.exception correctly to print tracebacks --- packages/jupyter-ai/jupyter_ai/personas/persona_manager.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index 12645939c..ca2893c58 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -373,9 +373,7 @@ def load_from_dir(root_dir: str, log: Logger) -> list[type[BasePersona]]: except Exception as e: # On exception, log error and continue to next file - # This mirrors the error handling pattern from entry point loading - log.error(f"Unable to load persona classes from '{py_file}'") - log.error(f"Error was: {type(e).__name__}: {e} - Unable to load persona classes from '{py_file}'") + log.exception(f"Unable to load persona classes from '{py_file}', exception details printed below.") continue return persona_classes From 5325ba6fb774f7db7d1fc573485a67abd064a77b Mon Sep 17 00:00:00 2001 From: Fernando Perez Date: Thu, 26 Jun 2025 23:35:03 -0700 Subject: [PATCH 12/13] Fix path loading logic so personas can import local resources --- .../jupyter_ai/personas/persona_manager.py | 69 +++++++++++-------- 1 file changed, 40 insertions(+), 29 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index ca2893c58..6a125dca0 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -4,6 +4,7 @@ import importlib.util import inspect import os +import sys from glob import glob from logging import Logger from pathlib import Path @@ -345,36 +346,46 @@ def load_from_dir(root_dir: str, log: Logger) -> list[type[BasePersona]]: if py_files: log.info(f"Found files from {root_dir}: {[Path(f).name for f in py_files]}") - # For each .py file, dynamically import the module and extract all - # BasePersona subclasses. - for py_file in py_files: - try: - # Get module name from file path - module_name = Path(py_file).stem - - # Create module spec and load the module - spec = importlib.util.spec_from_file_location(module_name, py_file) - if spec is None or spec.loader is None: + # Temporarily add root_dir to sys.path for imports + root_dir_in_path = root_dir in sys.path + if not root_dir_in_path: + sys.path.insert(0, root_dir) + + try: + # For each .py file, dynamically import the module and extract all + # BasePersona subclasses. + for py_file in py_files: + try: + # Get module name from file path + module_name = Path(py_file).stem + + # Create module spec and load the module + spec = importlib.util.spec_from_file_location(module_name, py_file) + if spec is None or spec.loader is None: + continue + + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + # Find all classes in the module that are BasePersona subclasses + for name, obj in inspect.getmembers(module, inspect.isclass): + # Check if it's a subclass of BasePersona but not BasePersona itself + if ( + issubclass(obj, BasePersona) + and obj is not BasePersona + and obj.__module__ == module_name + ): + log.info(f"Found persona class '{obj.__name__}' in '{py_file}'") + persona_classes.append(obj) + + except Exception as e: + # On exception, log error and continue to next file + log.exception(f"Unable to load persona classes from '{py_file}', exception details printed below.") continue - - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - - # Find all classes in the module that are BasePersona subclasses - for name, obj in inspect.getmembers(module, inspect.isclass): - # Check if it's a subclass of BasePersona but not BasePersona itself - if ( - issubclass(obj, BasePersona) - and obj is not BasePersona - and obj.__module__ == module_name - ): - log.info(f"Found persona class '{obj.__name__}' in '{py_file}'") - persona_classes.append(obj) - - except Exception as e: - # On exception, log error and continue to next file - log.exception(f"Unable to load persona classes from '{py_file}', exception details printed below.") - continue + finally: + # Remove root_dir from sys.path if we added it + if not root_dir_in_path and root_dir in sys.path: + sys.path.remove(root_dir) return persona_classes From 37c2203a811c373e93a7aa8bb59073b38950d734 Mon Sep 17 00:00:00 2001 From: "David L. Qiu" Date: Fri, 27 Jun 2025 10:57:02 -0700 Subject: [PATCH 13/13] fix mypy errors --- .../jupyter_ai/personas/persona_manager.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index 6a125dca0..81c3c8dd7 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -52,7 +52,11 @@ class PersonaManager(LoggingConfigurable): type for type checkers. """ - _persona_classes: ClassVar[list[type[BasePersona]] | None] = None + # TODO: the Persona classes from entry points should be stored as a class + # attribute, since they will not change at runtime. + # That should be injected into this instance attribute when personas defined + # under `.jupyter` are loaded. + _persona_classes: list[type[BasePersona]] | None = None _personas: dict[str, BasePersona] file_id: str @@ -152,7 +156,11 @@ def _init_persona_classes(self) -> None: ) # Load persona classes from local filesystem - persona_classes.extend(load_from_dir(self.get_dotjupyter_dir(), self.log)) + dotjupyter_dir = self.get_dotjupyter_dir() + if dotjupyter_dir is None: + self.log.info("No .jupyter directory found for loading local personas.") + else: + persona_classes.extend(load_from_dir(dotjupyter_dir, self.log)) self._persona_classes = persona_classes