Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 14 additions & 3 deletions jupyter_ai_persona_manager/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,20 @@ class PersonaManagerExtension(ExtensionApp):
def initialize(self, argv: Any = None) -> None:
super().initialize()

self.persona_managers_by_room: dict[str, PersonaManager] = {}
"""Cache of PersonaManager instances, indexed by room ID."""
@property
def persona_managers_by_room(self) -> dict[str, PersonaManager]:
"""
Dictionary of PersonaManager instances indexed by room ID.

This is accessible to other extensions via
`self.settings['jupyter-ai']['persona-managers']`.
"""
if 'jupyter-ai' not in self.settings:
self.settings['jupyter-ai'] = {}
if 'persona-managers' not in self.settings['jupyter-ai']:
self.settings['jupyter-ai']['persona-managers'] = {}

return self.settings['jupyter-ai']['persona-managers']

@property
def event_loop(self) -> AbstractEventLoop:
Expand Down Expand Up @@ -114,7 +126,6 @@ def _on_router_chat_init(self, room_id: str, ychat: "YChat") -> None:

# Register persona manager callbacks with router
self.router.observe_chat_msg(room_id, persona_manager.on_chat_message)
self.router.observe_slash_cmd_msg(room_id, persona_manager.on_slash_cmd_message)

def _init_persona_manager(
self, room_id: str, ychat: "YChat"
Expand Down
43 changes: 6 additions & 37 deletions jupyter_ai_persona_manager/persona_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -416,46 +416,15 @@ def _broadcast(
self.event_loop.create_task(persona.process_message(message))
return

def on_slash_cmd_message(self, room_id: str, message: Message):
async def refresh_personas(self):
"""
Routes & handles a message containing a slash command. Returns `True` if
the message specified a valid slash command recognized by
`PersonaManager`, `False` otherwise. Notes:
Method that reloads all persona classes defined locally under
`.jupyter/personas`, and re-initializes each persona class available in
the current chat.

- Each message may have exactly one slash command, which must be
specified by the first word of the message.

- This method will return `True` even if the command was not handled
successfully. `False` is just meant to indicate that the control flow
should return back to `route_message()`. This allows AI personas to
receive custom slash commands that only they recognize.
"""
first_word = get_first_word(message.body)
assert first_word and first_word.startswith("/")

command_id = first_word[1:]
if command_id == "refresh-personas":
self.handle_refresh_personas_command(message)
return True

# If command is unrecognized, log an error
self.log.warning(f"Unrecognized slash command: '/{command_id}'")
return False

def handle_refresh_personas_command(self, _: Message) -> None:
This method is public because it is called by `jupyter_ai_chat_commands`
when the `/refresh-personas` slash command is sent.
"""
Handles the '/refresh-personas' slash command.

TODO: How do we show status/completion in the UI?
"""
self.log.info(
f"Received '/refresh-personas'. Refreshing personas in chat '{self.room_id}'..."
)

# Refresh personas in background task
asyncio.create_task(self._refresh_personas())

async def _refresh_personas(self):
# Shutdown all personas
await self.shutdown_personas()

Expand Down
Loading