Add avatar path support for personas

README.md

@@ -22,14 +22,19 @@ To create and register a custom AI persona:
 ```python
 from jupyter_ai_persona_manager import BasePersona, PersonaDefaults
 from jupyterlab_chat.models import Message
+import os
+
+# Path to avatar file in your package
+AVATAR_PATH = os.path.join(os.path.dirname(__file__), "assets", "avatar.svg")
+
 
 class MyCustomPersona(BasePersona):
     @property
     def defaults(self):
         return PersonaDefaults(
             name="MyPersona",
             description="A helpful custom assistant",
-            avatar_path="/api/ai/static/custom-avatar.svg",
+            avatar_path=AVATAR_PATH,  # Absolute path to avatar file
             system_prompt="You are a helpful assistant specialized in...",
         )
 
@@ -39,6 +44,8 @@ class MyCustomPersona(BasePersona):
         self.send_message(response)
 ```
 
+**Avatar Path**: The `avatar_path` should be an absolute path to an image file (SVG, PNG, or JPG) within your package. The avatar will be automatically served at `/api/ai/avatars/{filename}`. If multiple personas use the same filename, the first one found will be served.
+
 ### 2. Register via Entry Points
 
 Add to your package's `pyproject.toml`:
@@ -85,21 +92,28 @@ For development and local customization, personas can be loaded from the `.jupyt
 ```python
 from jupyter_ai_persona_manager import BasePersona, PersonaDefaults
 from jupyterlab_chat.models import Message
+import os
+
+# Path to avatar file (in same directory as persona file)
+AVATAR_PATH = os.path.join(os.path.dirname(__file__), "avatar.svg")
+
 
 class MyLocalPersona(BasePersona):
     @property
     def defaults(self):
         return PersonaDefaults(
             name="Local Dev Assistant",
             description="A persona for local development",
-            avatar_path="/api/ai/static/jupyternaut.svg",
+            avatar_path=AVATAR_PATH,
             system_prompt="You help with local development tasks.",
         )
 
     async def process_message(self, message: Message):
         self.send_message(f"Local persona received: {message.body}")
 ```
 
+**Note**: Place your avatar file (e.g., `avatar.svg`) in the same directory as your persona file.
+
 ### Refreshing Personas
 
 Use the `/refresh-personas` slash command in any chat to reload personas without restarting JupyterLab:

jupyter_ai_persona_manager/base_persona.py

@@ -34,7 +34,7 @@ class PersonaDefaults(BaseModel):
     ################################################
     name: str  # e.g. "Jupyternaut"
     description: str  # e.g. "..."
-    avatar_path: str  # e.g. /avatars/jupyternaut.svg
+    avatar_path: str  # Absolute filesystem path to avatar image file (SVG, PNG, or JPG)
     system_prompt: str  # e.g. "You are a language model named..."
 
     ################################################
@@ -171,15 +171,22 @@ def name(self) -> str:
     @property
     def avatar_path(self) -> str:
         """
-        Returns the URL route that serves the avatar shown on messages from this
-        persona in the chat. This sets the `avatar_url` field in the data model
-        returned by `self.as_user()`. Provided by `BasePersona`.
+        Returns the API URL route that serves the avatar for this persona.
 
-        NOTE/TODO: This currently just returns the value set in `self.defaults`.
-        This is set here because we may require this field to be configurable
-        for all personas in the future.
+        The avatar is served at `/api/ai/avatars/{id}` where the ID is the
+        unique persona identifier. This ensures that each persona has a unique
+        avatar URL without exposing filesystem paths.
+
+        The actual avatar file path is specified in `defaults.avatar_path` as an
+        absolute filesystem path to an image file (SVG, PNG, or JPG) within the
+        persona's package or module.
+
+        This sets the `avatar_url` field in the data model returned by
+        `self.as_user()`. Provided by `BasePersona`.
         """
-        return self.defaults.avatar_path
+        # URL-encode the persona ID to handle special characters
+        from urllib.parse import quote
+        return f"/api/ai/avatars/{quote(self.id, safe='')}"
 
     @property
     def system_prompt(self) -> str:

jupyter_ai_persona_manager/extension.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import time
+from asyncio import get_event_loop_policy
 from typing import TYPE_CHECKING
 
 from jupyter_server.extension.application import ExtensionApp
@@ -10,7 +11,7 @@
 from traitlets import Type
 from traitlets.config import Config
 
-from jupyter_ai_persona_manager.handlers import RouteHandler
+from jupyter_ai_persona_manager.handlers import AvatarHandler, build_avatar_cache
 
 from .persona_manager import PersonaManager
 
@@ -30,8 +31,8 @@ class PersonaManagerExtension(ExtensionApp):
 
     name = "jupyter_ai_persona_manager"
     handlers = [
-        (r"jupyter-ai-persona-manager/health/?", RouteHandler)
-    ]  # No direct HTTP handlers, works through router integration
+        (r"/api/ai/avatars/(.*)", AvatarHandler),
+    ]
 
     persona_manager_class = Type(
         klass=PersonaManager,
@@ -48,28 +49,27 @@ def event_loop(self) -> AbstractEventLoop:
         """
         Returns a reference to the asyncio event loop.
         """
-        from asyncio import get_event_loop_policy
         return get_event_loop_policy().get_event_loop()
 
     def initialize_settings(self):
         """Initialize persona manager settings and router integration."""
         start = time.time()
-        
+
         # Ensure 'jupyter-ai.persona-manager' is in `self.settings`, which gets
         # copied to `self.serverapp.web_app.settings` after this method returns
         if 'jupyter-ai' not in self.settings:
             self.settings['jupyter-ai'] = {}
         if 'persona-manager' not in self.settings['jupyter-ai']:
             self.settings['jupyter-ai']['persona-managers'] = {}
-        
+
         # Set up router integration task
         self.event_loop.create_task(self._setup_router_integration())
 
         # Log server extension startup time
         self.log.info(f"Registered {self.name} server extension")
         startup_time = round((time.time() - start) * 1000)
         self.log.info(f"Initialized Persona Manager server extension in {startup_time} ms.")
-    
+
     async def _setup_router_integration(self) -> None:
         """
         Set up integration with jupyter-ai-router.
@@ -110,7 +110,7 @@ def _on_router_chat_init(self, room_id: str, ychat: "YChat") -> None:
         This initializes persona manager for the new chat room.
         """
         self.log.info(f"Router detected new chat room, initializing persona manager: {room_id}")
-        
+
         # Initialize persona manager for this chat
         persona_manager = self._init_persona_manager(room_id, ychat)
         if not persona_manager:
@@ -119,7 +119,7 @@ def _on_router_chat_init(self, room_id: str, ychat: "YChat") -> None:
                 + "Please verify your configuration and open a new issue on GitHub if this error persists."
             )
             return
-        
+
         # Cache the persona manager in server settings dictionary.
         #
         # NOTE: This must be added to `self.serverapp.web_app.settings`, not
@@ -128,7 +128,10 @@ def _on_router_chat_init(self, room_id: str, ychat: "YChat") -> None:
         # `self.initialize_settings` returns.
         persona_managers_by_room = self.serverapp.web_app.settings['jupyter-ai']['persona-managers']
         persona_managers_by_room[room_id] = persona_manager
-
+
+        # Rebuild avatar cache to include the new personas
+        build_avatar_cache(persona_managers_by_room)
+
         # Register persona manager callbacks with router
         self.router.observe_chat_msg(room_id, persona_manager.on_chat_message)
 

jupyter_ai_persona_manager/handlers.py

@@ -1,14 +1,99 @@
 import json
+import mimetypes
+import os
+from typing import Dict, Optional
+from urllib.parse import unquote
 
-from jupyter_server.base.handlers import APIHandler
+from jupyter_server.base.handlers import JupyterHandler
 import tornado
 
-class RouteHandler(APIHandler):
-    # The following decorator should be present on all verb methods (head, get, post,
-    # patch, put, delete, options) to ensure only authorized user can request the
-    # Jupyter server
+
+# Maximum avatar file size (5MB)
+MAX_AVATAR_SIZE = 5 * 1024 * 1024
+
+# Module-level cache: {persona_id: avatar_path}
+# This is populated when personas are initialized/refreshed
+_avatar_cache: Dict[str, str] = {}
+
+
+def build_avatar_cache(persona_managers: dict) -> None:
+    """
+    Build the avatar cache from all persona managers.
+
+    This should be called when personas are initialized or refreshed.
+    """
+    global _avatar_cache
+    _avatar_cache = {}
+
+    for room_id, persona_manager in persona_managers.items():
+        for persona in persona_manager.personas.values():
+            try:
+                avatar_path = persona.defaults.avatar_path
+                if avatar_path and os.path.exists(avatar_path):
+                    _avatar_cache[persona.id] = avatar_path
+            except Exception:
+                # Skip personas with invalid avatar paths
+                continue
+
+
+def clear_avatar_cache() -> None:
+    """Clear the avatar cache. Called during persona refresh."""
+    global _avatar_cache
+    _avatar_cache = {}
+
+
+class AvatarHandler(JupyterHandler):
+    """
+    Handler for serving persona avatar files.
+
+    Looks up avatar files by persona ID and serves the image file
+    with appropriate content-type headers.
+    """
+
     @tornado.web.authenticated
-    def get(self):
-        self.finish(json.dumps({
-            "data": "This is /jupyter-ai-persona-manager/get-example endpoint!"
-        }))
+    async def get(self, persona_id: str):
+        """Serve an avatar file by persona ID."""
+        # URL-decode the persona ID
+        persona_id = unquote(persona_id)
+
+        # Get the avatar file path
+        avatar_path = self._find_avatar_file(persona_id)
+
+        if avatar_path is None:
+            raise tornado.web.HTTPError(404, f"Avatar not found for persona")
+
+        # Check file size
+        try:
+            file_size = os.path.getsize(avatar_path)
+            if file_size > MAX_AVATAR_SIZE:
+                self.log.error(f"Avatar file too large: {file_size} bytes (max: {MAX_AVATAR_SIZE})")
+                raise tornado.web.HTTPError(413, "Avatar file too large")
+        except OSError as e:
+            self.log.error(f"Error checking avatar file size: {e}")
+            raise tornado.web.HTTPError(500, "Error accessing avatar file")
+
+        # Serve the file
+        try:
+            # Set content type based on file extension
+            content_type, _ = mimetypes.guess_type(avatar_path)
+            if content_type:
+                self.set_header("Content-Type", content_type)
+
+            # Read and serve the file
+            with open(avatar_path, 'rb') as f:
+                content = f.read()
+                self.write(content)
+
+            await self.finish()
+        except Exception as e:
+            self.log.error(f"Error serving avatar file: {e}")
+            raise tornado.web.HTTPError(500, f"Error serving avatar file: {str(e)}")
+
+    def _find_avatar_file(self, persona_id: str) -> Optional[str]:
+        """
+        Find the avatar file path by persona ID using the module-level cache.
+
+        The cache is built when personas are initialized or refreshed,
+        so this is an O(1) lookup instead of iterating all personas.
+        """
+        return _avatar_cache.get(persona_id)

jupyter_ai_persona_manager/persona_manager.py

@@ -19,6 +19,7 @@
 
 from .base_persona import BasePersona
 from .directories import find_dot_dir, find_workspace_dir
+from .handlers import build_avatar_cache
 
 if TYPE_CHECKING:
     from asyncio import AbstractEventLoop
@@ -284,6 +285,7 @@ def _init_personas(self) -> dict[str, BasePersona]:
         self.log.info(
             f"SUCCESS: Initialized {len(personas)} AI personas for chat room '{self.ychat.get_id()}'. Time elapsed: {elapsed_time_ms}ms."
         )
+
         return personas
 
     def _display_persona_error_message(self, persona_item: dict) -> None:
@@ -432,6 +434,16 @@ async def refresh_personas(self):
         self._init_local_persona_classes()
         self._personas = self._init_personas()
 
+        # Rebuild avatar cache after reloading personas
+        # Get all persona managers from parent (extension) settings
+        try:
+            # Access all persona managers through the parent extension's settings
+            # Note: self.parent is the PersonaManagerExtension instance
+            persona_managers = self.parent.serverapp.web_app.settings.get('jupyter-ai', {}).get('persona-managers', {})
+            build_avatar_cache(persona_managers)
+        except Exception as e:
+            self.log.error(f"Error rebuilding avatar cache: {e}")
+
         # Write success message to chat & logs
         self.send_system_message("Refreshed all AI personas in this chat.")
         self.log.info(f"Refreshed all AI personas in chat '{self.room_id}'.")