Add avatar path support for personas

This adds the ability for personas to specify an avatar image file directly in their package instead of needing to implement their own static file handler.

Personas now set `avatar_path` to an absolute path pointing to an image file (SVG, PNG, JPG). The server dynamically looks up and serves these files through `/api/ai/avatars/{url_encoded_id}`.

- Uses persona.id (URL-encoded) for avatar URLs to ensure uniqueness
- File size validation (5MB max)
- Proper content-type headers based on file extension
- Module-level cache for O(1) avatar path lookups (built at init, rebuilt on refresh)

Fixes #6

Author: Zsailer

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. "/path/to/package/avatars/jupyternaut.svg" - absolute path to avatar file
+    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,16 +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`.
         """
-        filename = os.path.basename(self.defaults.avatar_path)
-        return f"/api/ai/avatars/{filename}"
+        # 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

@@ -11,7 +11,7 @@
 from traitlets import Type
 from traitlets.config import Config
 
-from jupyter_ai_persona_manager.handlers import RouteHandler, AvatarHandler
+from jupyter_ai_persona_manager.handlers import AvatarHandler, build_avatar_cache
 
 from .persona_manager import PersonaManager
 
@@ -31,7 +31,6 @@ class PersonaManagerExtension(ExtensionApp):
 
     name = "jupyter_ai_persona_manager"
     handlers = [
-        (r"jupyter-ai-persona-manager/health/?", RouteHandler),
         (r"/api/ai/avatars/(.*)", AvatarHandler),
     ]
 
@@ -130,6 +129,9 @@ def _on_router_chat_init(self, room_id: str, ychat: "YChat") -> None:
         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,41 +1,76 @@
 import json
 import mimetypes
 import os
-from functools import lru_cache
-from typing import Optional
+from typing import Dict, Optional
+from urllib.parse import unquote
 
 from jupyter_server.base.handlers import JupyterHandler
 import tornado
 
 
-class RouteHandler(JupyterHandler):
-    # 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
-    @tornado.web.authenticated
-    def get(self):
-        self.set_header("Content-Type", "application/json")
-        self.finish(json.dumps({
-            "data": "This is /jupyter-ai-persona-manager/get-example endpoint!"
-        }))
+# 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 through the PersonaManager to find the correct file path,
-    then serves the file with appropriate content-type headers.
+    Looks up avatar files by persona ID and serves the image file
+    with appropriate content-type headers.
     """
 
     @tornado.web.authenticated
-    async def get(self, filename: str):
-        """Serve an avatar file by filename."""
-        # Get the avatar file path from persona managers
-        avatar_path = self._find_avatar_file(filename)
+    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 file not found: {filename}")
+            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:
@@ -51,32 +86,14 @@ async def get(self, filename: str):
 
             await self.finish()
         except Exception as e:
-            self.log.error(f"Error serving avatar file {filename}: {e}")
+            self.log.error(f"Error serving avatar file: {e}")
             raise tornado.web.HTTPError(500, f"Error serving avatar file: {str(e)}")
 
-    @lru_cache(maxsize=128)
-    def _find_avatar_file(self, filename: str) -> Optional[str]:
+    def _find_avatar_file(self, persona_id: str) -> Optional[str]:
         """
-        Find the avatar file path by searching through all persona managers.
+        Find the avatar file path by persona ID using the module-level cache.
 
-        Uses LRU cache to avoid repeated lookups for the same filename.
+        The cache is built when personas are initialized or refreshed,
+        so this is an O(1) lookup instead of iterating all personas.
         """
-        # Get all persona managers from settings
-        persona_managers = self.settings.get('jupyter-ai', {}).get('persona-managers', {})
-
-        for room_id, persona_manager in persona_managers.items():
-            # Check each persona's avatar path
-            for persona in persona_manager.personas.values():
-                try:
-                    avatar_path = persona.defaults.avatar_path
-                    if avatar_path and os.path.basename(avatar_path) == filename:
-                        # Found a match, return the absolute path
-                        if os.path.exists(avatar_path):
-                            return avatar_path
-                        else:
-                            self.log.warning(f"Avatar file not found at path: {avatar_path}")
-                except Exception as e:
-                    self.log.warning(f"Error checking avatar for persona {persona.name}: {e}")
-                    continue
-
-        return None
+        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
@@ -433,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}'.")

jupyter_ai_persona_manager/tests/test_handlers.py

@@ -3,39 +3,16 @@
 import tempfile
 from pathlib import Path
 from unittest.mock import Mock
+from urllib.parse import quote
 
 import pytest
 
-
-async def test_health(jp_fetch):
-    # When
-    response = await jp_fetch("jupyter-ai-persona-manager", "health")
-
-    # Then
-    assert response.code == 200
-    payload = json.loads(response.body)
-    assert payload == {
-        "data": "This is /jupyter-ai-persona-manager/get-example endpoint!"
-    }
-
-
-@pytest.fixture
-def mock_persona_with_avatar(tmp_path):
-    """Create a mock persona with an avatar file."""
-    # Create avatar file
-    avatar_file = tmp_path / "test_avatar.svg"
-    avatar_file.write_text('<svg><circle r="10"/></svg>')
-
-    # Create mock persona
-    mock_persona = Mock()
-    mock_persona.defaults.avatar_path = str(avatar_file)
-    mock_persona.name = "TestPersona"
-
-    return mock_persona, str(avatar_file)
+from jupyter_ai_persona_manager.handlers import build_avatar_cache
 
 
 async def test_avatar_handler_serves_file(jp_fetch, jp_serverapp, tmp_path):
     """Test that the avatar handler can serve avatar files."""
+
     # Create avatar file
     avatar_file = tmp_path / "test.svg"
     avatar_file.write_text('<svg><circle r="10"/></svg>')
@@ -44,6 +21,7 @@ async def test_avatar_handler_serves_file(jp_fetch, jp_serverapp, tmp_path):
     mock_persona = Mock()
     mock_persona.defaults.avatar_path = str(avatar_file)
     mock_persona.name = "TestPersona"
+    mock_persona.id = "jupyter-ai-personas::test::TestPersona"
 
     # Create mock persona manager
     mock_pm = Mock()
@@ -56,8 +34,12 @@ async def test_avatar_handler_serves_file(jp_fetch, jp_serverapp, tmp_path):
         'room1': mock_pm
     }
 
-    # Fetch the avatar
-    response = await jp_fetch("api", "ai", "avatars", "test.svg")
+    # Build the avatar cache
+    build_avatar_cache(jp_serverapp.web_app.settings['jupyter-ai']['persona-managers'])
+
+    # Fetch the avatar using URL-encoded persona ID
+    encoded_id = quote(mock_persona.id, safe='')
+    response = await jp_fetch("api", "ai", "avatars", encoded_id)
 
     # Verify response
     assert response.code == 200
@@ -67,6 +49,7 @@ async def test_avatar_handler_serves_file(jp_fetch, jp_serverapp, tmp_path):
 
 async def test_avatar_handler_404_for_missing_file(jp_fetch, jp_serverapp):
     """Test that the avatar handler returns 404 for missing files."""
+
     # Create mock persona manager with no matching avatar
     mock_pm = Mock()
     mock_pm.personas = {}
@@ -78,16 +61,20 @@ async def test_avatar_handler_404_for_missing_file(jp_fetch, jp_serverapp):
         'room1': mock_pm
     }
 
+    # Build the avatar cache (will be empty)
+    build_avatar_cache(jp_serverapp.web_app.settings['jupyter-ai']['persona-managers'])
+
     # Try to fetch a non-existent avatar
     with pytest.raises(Exception) as exc_info:
-        await jp_fetch("api", "ai", "avatars", "nonexistent.svg")
+        await jp_fetch("api", "ai", "avatars", "nonexistent-id")
 
     # Verify 404 response
     assert '404' in str(exc_info.value) or 'Not Found' in str(exc_info.value)
 
 
 async def test_avatar_handler_serves_png(jp_fetch, jp_serverapp, tmp_path):
     """Test that the avatar handler can serve PNG files."""
+
     # Create PNG file
     avatar_file = tmp_path / "test.png"
     avatar_file.write_bytes(b'\x89PNG\r\n\x1a\n')
@@ -96,6 +83,7 @@ async def test_avatar_handler_serves_png(jp_fetch, jp_serverapp, tmp_path):
     mock_persona = Mock()
     mock_persona.defaults.avatar_path = str(avatar_file)
     mock_persona.name = "TestPersona"
+    mock_persona.id = "jupyter-ai-personas::test::AnotherPersona"
 
     # Create mock persona manager
     mock_pm = Mock()
@@ -108,10 +96,16 @@ async def test_avatar_handler_serves_png(jp_fetch, jp_serverapp, tmp_path):
         'room1': mock_pm
     }
 
-    # Fetch the avatar
-    response = await jp_fetch("api", "ai", "avatars", "test.png")
+    # Build the avatar cache
+    build_avatar_cache(jp_serverapp.web_app.settings['jupyter-ai']['persona-managers'])
+
+    # Fetch the avatar using URL-encoded persona ID
+    encoded_id = quote(mock_persona.id, safe='')
+    response = await jp_fetch("api", "ai", "avatars", encoded_id)
 
     # Verify response
     assert response.code == 200
     assert response.body.startswith(b'\x89PNG')
     assert 'image/png' in response.headers.get('Content-Type', '')
+
+