diff --git a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py index 039da0e15..0fb579cda 100644 --- a/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py +++ b/packages/jupyter-ai/jupyter_ai/personas/persona_manager.py @@ -209,8 +209,21 @@ def _init_local_persona_classes(self) -> None: dotjupyter_dir = self.get_dotjupyter_dir() if dotjupyter_dir is None: self.log.info("No .jupyter directory found for loading local personas.") - else: - self._local_persona_classes = load_from_dir(dotjupyter_dir, self.log) + return + + if find_persona_files(dotjupyter_dir): + self.send_system_message( + "Found persona files in `.jupyter` directory. Please move them to `.jupyter/personas/` subdirectory." + ) + + personas_subdir = os.path.join(dotjupyter_dir, "personas") + if not os.path.exists(personas_subdir): + self.log.info( + "No `personas` subdirectory found in `.jupyter` directory for loading local personas." + ) + return + + self._local_persona_classes = load_from_dir(personas_subdir, self.log) def _init_personas(self) -> dict[str, BasePersona]: """ @@ -562,6 +575,25 @@ def is_persona(username: str): return username.startswith("jupyter-ai-personas") +def find_persona_files(dir: str) -> list[str]: + """Find persona Python files in a directory without loading them.""" + if not os.path.exists(dir): + return [] + + try: + all_py_files = glob(os.path.join(dir, "*.py")) + 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) + return py_files + except Exception: + return [] + + def load_from_dir(dir: str, log: Logger) -> list[dict]: """ Load _persona class declarations_ from Python files in the local filesystem. @@ -583,31 +615,11 @@ def load_from_dir(dir: str, log: Logger) -> list[dict]: """ persona_classes: list[dict] = [] - log.info(f"Searching for persona files in {dir}") - # Check if root directory exists - if not os.path.exists(dir): - return persona_classes - - # Find all .py files in the root directory that contain "persona" in the name - try: - all_py_files = glob(os.path.join(dir, "*.py")) - 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.error( - f"{type(e).__name__} occurred while searching for Python files in {dir}" - ) + py_files = find_persona_files(dir) + if not py_files: return persona_classes - if py_files: - log.info(f"Found files from {dir}: {[Path(f).name for f in py_files]}") + log.info(f"Loading persona files from {dir}: {[Path(f).name for f in py_files]}") # Temporarily add root_dir to sys.path for imports dir_in_path = dir in sys.path diff --git a/packages/jupyter-ai/jupyter_ai/tests/test_personas.py b/packages/jupyter-ai/jupyter_ai/tests/test_personas.py index 8256d02df..886ac2164 100644 --- a/packages/jupyter-ai/jupyter_ai/tests/test_personas.py +++ b/packages/jupyter-ai/jupyter_ai/tests/test_personas.py @@ -7,8 +7,8 @@ from unittest.mock import Mock import pytest -from jupyter_ai.personas.base_persona import BasePersona, PersonaDefaults -from jupyter_ai.personas.persona_manager import load_from_dir +from jupyter_ai.personas.base_persona import BasePersona +from jupyter_ai.personas.persona_manager import find_persona_files, load_from_dir @pytest.fixture @@ -24,6 +24,39 @@ def mock_logger(): return Mock() +class TestFindPersonaFiles: + """Test cases for find_persona_files function.""" + + def test_nonexistent_directory_returns_empty_list(self): + """Test that a non-existent directory returns an empty list.""" + result = find_persona_files("/nonexistent/directory/path") + assert result == [] + + def test_empty_directory_returns_empty_list(self, tmp_persona_dir): + """Test that an empty directory returns an empty list.""" + result = find_persona_files(str(tmp_persona_dir)) + assert result == [] + + def test_finds_valid_ignores_invalid_persona_files(self, tmp_persona_dir): + """Test that persona files with valid filenames are found while private, hidden, and non-valid files are ignored.""" + (tmp_persona_dir / "my_persona.py").write_text("pass") + (tmp_persona_dir / "PersonalAssistant.py").write_text("pass") + + (tmp_persona_dir / "my_other_code.py").write_text("pass") + (tmp_persona_dir / "_private_persona.py").write_text("pass") + (tmp_persona_dir / ".hidden_persona.py").write_text("pass") + + result = find_persona_files(str(tmp_persona_dir)) + result_names = [Path(f).name for f in result] + + assert "my_persona.py" in result_names + assert "PersonalAssistant.py" in result_names + + assert "my_other_code.py" not in result_names + assert "_private_persona.py" not in result_names + assert ".hidden_persona.py" not in result_names + + class TestLoadPersonaClassesFromDirectory: """Test cases for load_from_dir function."""