Skip to content

Commit c464c17

Browse files
committed
First cut at the PersonaLoader imnplementation, no tests yet.
1 parent b597e5c commit c464c17

File tree

1 file changed

+112
-0
lines changed

1 file changed

+112
-0
lines changed

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

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
11
from __future__ import annotations
22

33
import asyncio
4+
import importlib.util
5+
import inspect
46
import os
7+
from glob import glob
58
from logging import Logger
9+
from pathlib import Path
610
from time import time_ns
711
from typing import TYPE_CHECKING, Any, ClassVar
812

@@ -140,6 +144,10 @@ def _init_persona_classes(self) -> None:
140144
"ERROR: Jupyter AI has no AI personas available. "
141145
+ "Please verify your server configuration and open a new issue on our GitHub repo if this warning persists."
142146
)
147+
148+
# 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.
149+
persona_classes.extend(LocalPersonaLoader(self.root_dir).load_persona_classes())
150+
143151
PersonaManager._persona_classes = persona_classes
144152

145153
def _init_personas(self) -> dict[str, BasePersona]:
@@ -287,3 +295,107 @@ def get_mcp_config(self) -> dict[str, Any]:
287295
return {}
288296
else:
289297
return self._mcp_config_loader.get_config(jdir)
298+
299+
300+
class LocalPersonaLoader(LoggingConfigurable):
301+
"""
302+
Load _persona class declarations_ from the local filesystem.
303+
304+
Those class declarations are then used to instantiate personas by the `PersonaManager`.
305+
306+
TODO: wire to local .jupyter directory system, for now just use the current
307+
directory.
308+
"""
309+
310+
def __init__(
311+
self,
312+
*args,
313+
root_dir: str | None = None,
314+
**kwargs,
315+
):
316+
# Forward other arguments to parent class
317+
super().__init__(*args, **kwargs)
318+
319+
# Set default root directory to current working directory if not provided
320+
self.root_dir = root_dir or os.getcwd()
321+
322+
self.log.info(f"LocalPersonaLoader initialized with root directory: {self.root_dir}")
323+
324+
def load_persona_classes(self) -> list[type[BasePersona]]:
325+
"""
326+
Loads persona classes from Python files in the local filesystem.
327+
328+
Scans the root_dir for .py files, dynamically imports them, and extracts
329+
any class declarations that are subclasses of BasePersona.
330+
"""
331+
persona_classes: list[type[BasePersona]] = []
332+
333+
# Check if root directory exists
334+
if not os.path.exists(self.root_dir):
335+
self.log.info(f"Root directory does not exist: {self.root_dir}")
336+
return persona_classes
337+
338+
# Find all .py files in the root directory
339+
py_files = glob(os.path.join(self.root_dir, "*.py"))
340+
341+
if not py_files:
342+
self.log.info(f"No Python files found in directory: {self.root_dir}")
343+
return persona_classes
344+
345+
self.log.info(f"Found {len(py_files)} Python files in {self.root_dir}")
346+
self.log.info("PENDING: Loading persona classes from local Python files...")
347+
start_time_ns = time_ns()
348+
349+
for py_file in py_files:
350+
try:
351+
# Get module name from file path
352+
module_name = Path(py_file).stem
353+
354+
# Skip if module name starts with underscore (private modules)
355+
if module_name.startswith('_'):
356+
continue
357+
358+
# Create module spec and load the module
359+
spec = importlib.util.spec_from_file_location(module_name, py_file)
360+
if spec is None or spec.loader is None:
361+
self.log.warning(f" - Unable to create module spec for {py_file}")
362+
continue
363+
364+
module = importlib.util.module_from_spec(spec)
365+
spec.loader.exec_module(module)
366+
367+
# Find all classes in the module that are BasePersona subclasses
368+
module_persona_classes = []
369+
for name, obj in inspect.getmembers(module, inspect.isclass):
370+
# Check if it's a subclass of BasePersona but not BasePersona itself
371+
if (issubclass(obj, BasePersona) and
372+
obj is not BasePersona and
373+
obj.__module__ == module_name):
374+
module_persona_classes.append(obj)
375+
376+
if module_persona_classes:
377+
persona_classes.extend(module_persona_classes)
378+
class_names = [cls.__name__ for cls in module_persona_classes]
379+
self.log.info(
380+
f" - Loaded {len(module_persona_classes)} persona class(es) from '{py_file}': {class_names}"
381+
)
382+
else:
383+
self.log.debug(f" - No persona classes found in '{py_file}'")
384+
385+
except Exception:
386+
# On exception, log an error and continue
387+
# This mirrors the error handling pattern from entry point loading
388+
self.log.exception(
389+
f" - Unable to load persona classes from '{py_file}' due to an exception printed below."
390+
)
391+
continue
392+
393+
if len(persona_classes) > 0:
394+
elapsed_time_ms = (time_ns() - start_time_ns) // 1_000_000
395+
self.log.info(
396+
f"SUCCESS: Loaded {len(persona_classes)} persona classes from local filesystem. Time elapsed: {elapsed_time_ms}ms."
397+
)
398+
else:
399+
self.log.info("No persona classes found in local filesystem.")
400+
401+
return persona_classes

0 commit comments

Comments
 (0)