"""
AI Assistant Components for Rubric AI
Contains: ai_assistant_pane, suggestion_card, chat_input
"""
from contextlib import contextmanager
from nicegui import ui
from .layout import material_icon
from frontend.utils.formatting import format_relative_timestamp
from frontend.config.stage_config import get_stage_config, get_stage_display_name
@contextmanager
def ai_assistant_pane():
"""
Create the AI assistant pane (right side of the two-column layout).
This pane contains:
- Header with AI icon and title
- Content area for suggestions/responses
- Chat input at the bottom
Usage:
with ai_assistant_pane() as assistant:
suggestion_card(
message="I've drafted 3 groups for you.",
suggestions=['Thesis', 'Evidence', 'Clarity'],
on_add_all=handle_add_all
)
"""
with ui.element('div').classes('pane p-6'):
# Header
with ui.element('div').classes('ai-assistant-header'):
with ui.element('div').classes('ai-icon-wrapper'):
material_icon('auto_awesome')
ui.html('
AI Assistant
', sanitize=False)
# Content area (will be filled by yielded content)
with ui.element('div').classes('flex-grow flex flex-col') as content:
yield content
class SuggestionCard:
"""
A card showing AI suggestions with action buttons.
Displays:
- Message from the AI
- Clickable suggestion chips
- "Add all to Workspace" button
- "Add Individually" button
"""
def __init__(
self,
message: str,
suggestions: list[str],
on_add_all=None,
on_add_individual=None,
on_chip_click=None
):
"""
Create a suggestion card.
Args:
message: AI message to display
suggestions: List of suggested items
on_add_all: Callback when "Add all" is clicked (receives list)
on_add_individual: Callback to show individual add UI
on_chip_click: Callback when a chip is clicked (receives chip text)
"""
self.message = message
self.suggestions = suggestions
self.on_add_all = on_add_all
self.on_add_individual = on_add_individual
self.on_chip_click = on_chip_click
self.element = None
self._build()
def _build(self):
"""Build the UI elements."""
with ui.element('div').classes('ai-suggestion-card') as card:
self.element = card
# Message
ui.html(f'{self.message}
', sanitize=False)
# Suggestion chips
with ui.element('div').classes('suggestion-chips'):
for suggestion in self.suggestions:
chip = ui.button(suggestion).classes('suggestion-chip')
if self.on_chip_click:
chip.on('click', lambda s=suggestion: self.on_chip_click(s))
# Action buttons
with ui.element('div').classes('ai-action-buttons'):
# Add all button (gradient)
ui.button(
f'Add all {len(self.suggestions)} to Workspace',
on_click=lambda: self.on_add_all(self.suggestions) if self.on_add_all else None
).classes('btn btn-primary')
# Add individually button
ui.button(
'Add Individually +',
on_click=self.on_add_individual
).classes('btn btn-outline')
def suggestion_card(
message: str,
suggestions: list[str],
on_add_all=None,
on_add_individual=None,
on_chip_click=None
) -> SuggestionCard:
"""
Factory function to create a suggestion card.
Args:
message: AI message to display
suggestions: List of suggested items
on_add_all: Callback when "Add all" is clicked
on_add_individual: Callback to show individual add UI
on_chip_click: Callback when a chip is clicked
Returns:
SuggestionCard instance
"""
return SuggestionCard(
message=message,
suggestions=suggestions,
on_add_all=on_add_all,
on_add_individual=on_add_individual,
on_chip_click=on_chip_click
)
def chat_input(placeholder: str = 'Ask your assistant...', on_send=None):
"""
Create a chat input bar at the bottom of the assistant pane.
Args:
placeholder: Placeholder text for the input
on_send: Async callback when message is sent (receives message text)
Returns:
The input element for reference
"""
with ui.element('div').classes('chat-input-wrapper'):
input_el = ui.input(placeholder=placeholder).classes('flex-grow').props('borderless dense')
async def handle_send():
message = input_el.value
if message and on_send:
input_el.set_value('')
await on_send(message)
ui.button(icon='send', on_click=handle_send).props('flat round dense').classes('text-primary')
input_el.on('keydown.enter', handle_send)
return input_el
class ChatHistory:
"""
Manages a scrollable chat history display with auto-scroll.
"""
def __init__(self, initial_message: str = None):
"""
Create a chat history container.
Args:
initial_message: Optional initial AI message to display
"""
self.container = None
self.messages = []
with ui.element('div').classes('ai-chat-scroll-container') as container:
self.container = container
if initial_message:
self.add_ai_message(initial_message)
def _scroll_to_bottom(self):
"""
Scroll the chat container to the bottom if user is already near bottom.
This respects manual scrolling - if user has scrolled up to read history,
we don't force them back to the bottom. Only auto-scroll if they're
already at or near the bottom (within 100px threshold).
"""
if self.container:
# Check if user is near bottom before scrolling
ui.run_javascript(f'''
const container = document.querySelector('.ai-chat-scroll-container');
if (container) {{
const scrollThreshold = 100; // pixels from bottom
const isNearBottom = container.scrollHeight - container.scrollTop - container.clientHeight < scrollThreshold;
// Only auto-scroll if user is already near bottom
if (isNearBottom) {{
container.scrollTop = container.scrollHeight;
}}
}}
''')
def add_ai_message(self, message: str):
"""Add an AI message to the history and scroll to bottom."""
with self.container:
with ui.element('div').classes('mb-3'):
with ui.element('div').classes('ai-suggestion-text text-left'):
ui.markdown(message)
# Auto-scroll to bottom after adding message (if user is at bottom)
self._scroll_to_bottom()
def add_user_message(self, message: str):
"""Add a user message to the history and scroll to bottom."""
with self.container:
with ui.element('div').classes('flex gap-2 mb-3 justify-end'):
with ui.element('div').classes('bg-primary text-white rounded-lg px-3 py-2 max-w-xs'):
ui.label(message).classes('text-sm')
# Auto-scroll to bottom after adding message (if user is at bottom)
self._scroll_to_bottom()
def add_spinner(self):
"""Add a thinking spinner. Returns the spinner element for later removal."""
with self.container:
spinner = ui.spinner(type='dots').classes('self-start')
# Scroll to show spinner (if user is at bottom)
self._scroll_to_bottom()
return spinner
def add_system_message(self, message: str):
"""Add a system notification message to the history and scroll to bottom."""
with self.container:
with ui.element('div').classes('flex justify-center my-3'):
with ui.element('div').classes(
'bg-slate-100 text-slate-600 rounded-full px-4 py-2 text-sm inline-flex items-center gap-2'
):
ui.label(message)
# Auto-scroll to bottom (if user is at bottom)
self._scroll_to_bottom()
def clear(self):
"""Clear all messages from the chat history."""
self.container.clear()
self.messages = []
class HybridAssistant:
"""
Hybrid AI assistant that combines suggestion cards with chat functionality.
Shows suggestion cards for structured responses, but allows free-form
chat input for custom requests.
"""
def __init__(
self,
initial_message: str = None,
on_chat_send=None
):
"""
Create a hybrid assistant panel.
Args:
initial_message: Initial AI greeting message
on_chat_send: Async callback when user sends a chat message
"""
self.on_chat_send = on_chat_send
self.suggestion_container = None
self.chat_history = None
self.input_el = None
self._build(initial_message)
def _build(self, initial_message: str):
"""Build the UI elements."""
with ui.element('div').classes('pane p-6 h-full flex flex-col'):
# Header
with ui.element('div').classes('ai-assistant-header'):
with ui.element('div').classes('ai-icon-wrapper'):
ui.html('auto_awesome', sanitize=False)
ui.html('AI Assistant
', sanitize=False)
# Suggestion/Response area
with ui.element('div').classes('flex-grow flex flex-col overflow-hidden') as container:
self.suggestion_container = container
# Initial message card
if initial_message:
with ui.element('div').classes('ai-suggestion-card flex-grow'):
ui.html(f'{initial_message}
', sanitize=False)
# Chat input at bottom
self.input_el = chat_input(
placeholder='Ask your assistant...',
on_send=self._handle_chat
)
async def _handle_chat(self, message: str):
"""Handle chat message from user."""
if self.on_chat_send:
await self.on_chat_send(message)
def show_suggestions(
self,
message: str,
suggestions: list[str],
on_add_all=None,
on_add_individual=None
):
"""
Display a suggestion card, replacing any previous content.
Args:
message: AI message
suggestions: List of suggestions
on_add_all: Callback when "Add all" is clicked
on_add_individual: Callback for individual add
"""
self.suggestion_container.clear()
with self.suggestion_container:
suggestion_card(
message=message,
suggestions=suggestions,
on_add_all=on_add_all,
on_add_individual=on_add_individual
)
def show_message(self, message: str):
"""
Display a simple message card.
Args:
message: Message to display
"""
self.suggestion_container.clear()
with self.suggestion_container:
with ui.element('div').classes('ai-suggestion-card flex-grow'):
ui.html(f'{message}
', sanitize=False)
def show_loading(self):
"""Show a loading spinner."""
self.suggestion_container.clear()
with self.suggestion_container:
with ui.element('div').classes('ai-suggestion-card flex-grow flex items-center justify-center'):
ui.spinner(type='dots', size='lg')
class ChatAssistant:
"""
Pure chat-based AI assistant with scrollable message history.
Provides a simple conversation interface where messages accumulate
in a scrollable container, supporting back-and-forth dialogue.
"""
def __init__(
self,
initial_message: str = None,
on_send=None
):
"""
Create a chat assistant panel.
Args:
initial_message: Initial AI greeting message
on_send: Async callback when user sends a message (receives message text)
"""
self.on_send = on_send
self.chat_history = None
self.input_el = None
self.current_spinner = None
self._build(initial_message)
def _build(self, initial_message: str):
"""Build the UI elements."""
with ui.element('div').classes('pane p-6 ai-assistant-pane'):
# Header
with ui.element('div').classes('ai-assistant-header'):
with ui.element('div').classes('ai-icon-wrapper'):
ui.html('auto_awesome', sanitize=False)
ui.html('AI Assistant
', sanitize=False)
# Chat history (scrollable)
self.chat_history = ChatHistory(initial_message=initial_message)
# Chat input at bottom
self.input_el = chat_input(
placeholder='Ask your assistant...',
on_send=self._handle_send
)
async def _handle_send(self, message: str):
"""Handle message from user."""
# Add user message to history
self.add_user_message(message)
# Call the callback if provided
if self.on_send:
await self.on_send(message)
def add_ai_message(self, message: str):
"""Add an AI message to the chat history."""
self.hide_loading()
self.chat_history.add_ai_message(message)
def add_user_message(self, message: str):
"""Add a user message to the chat history."""
self.chat_history.add_user_message(message)
def show_loading(self):
"""Show a loading spinner in the chat."""
if not self.current_spinner:
self.current_spinner = self.chat_history.add_spinner()
def hide_loading(self):
"""Hide the loading spinner."""
if self.current_spinner:
try:
self.current_spinner.delete()
except (ValueError, AttributeError):
# Spinner already removed or doesn't exist
pass
finally:
self.current_spinner = None
class TabbedAssistant:
"""
AI Assistant with tabbed interface.
Tabs:
- Assistant: Chat interface (existing ChatHistory + chat_input)
- Recent Conversations: List of past chats
- Configuration: Output mode selection and prompt customization
- Prompt Preview: Read-only view of full context
"""
def __init__(
self,
initial_message: str = None,
on_send=None,
conversations: list[dict] = None,
on_load_conversation=None,
on_delete_conversation=None,
on_new_chat=None,
stage_context: str = "stage1",
on_navigate=None,
):
"""
Create a tabbed assistant panel.
Args:
initial_message: Initial AI greeting message
on_send: Async callback when user sends a message
conversations: List of conversation dicts from Supabase
on_load_conversation: Async callback when user clicks a conversation
on_delete_conversation: Async callback when user deletes a conversation
on_new_chat: Callback when user clicks "Start New Chat"
stage_context: The current stage (stage1, stage2, etc.) - controls output modes
on_navigate: Callback when navigation to different stage is needed (receives stage, conversation_id)
"""
self.on_send = on_send
self.on_load_conversation = on_load_conversation
self.on_delete_conversation = on_delete_conversation
self.on_new_chat = on_new_chat
self.on_navigate = on_navigate
self.chat_history = None
self.input_el = None
self.current_spinner = None
self.stage_context = stage_context
# Configuration state
self.output_mode = 'standard'
self.mode_radio = None
self.mode_selector = None # Dropdown in Assistant tab header
self.prompt_description_label = None
self.system_prompt_input = None
self.prompt_preview_container = None
# Conversation state
self.current_conversation_id: str | None = None
self.conversations = conversations or []
self.conversations_list = None
self.tabs = None
self.assistant_tab = None
self._build(initial_message)
def _build(self, initial_message: str):
"""Build the tabbed UI."""
with ui.element('div').classes('pane p-6 ai-assistant-pane'):
# Header with mode selector
with ui.element('div').classes('ai-assistant-header'):
# Left side: icon and title
with ui.element('div').classes('flex items-center gap-3'):
with ui.element('div').classes('ai-icon-wrapper'):
ui.html('auto_awesome', sanitize=False)
ui.html('AI Assistant
', sanitize=False)
# Right side: mode selector
# Get stage-specific modes for dropdown options
stage_config = get_stage_config(self.stage_context)
modes = stage_config.get("modes", [])
mode_options = [m["display_name"] for m in modes]
default_mode = stage_config.get("default_mode", "standard")
# Find default display name
default_display = next(
(m["display_name"] for m in modes if m["key"] == default_mode),
mode_options[0] if mode_options else "Standard"
)
# Capsule/pill dropdown with icon
self.mode_selector = ui.select(
options=mode_options,
value=default_display,
on_change=self._handle_mode_selector_change
).props('dense borderless').classes('mode-selector-dropdown')
# Add sparkle icon to the left of the dropdown
self.mode_selector.props(add='prepend')
with self.mode_selector:
with self.mode_selector.add_slot('prepend'):
ui.html('auto_awesome', sanitize=False)
# Tabs
with ui.tabs().classes('ai-assistant-tabs') as tabs:
self.tabs = tabs
self.assistant_tab = ui.tab('Assistant')
conversations_tab = ui.tab('Recent Chats')
config_tab = ui.tab('Configuration')
preview_tab = ui.tab('Prompt Preview')
# Tab Panels
with ui.tab_panels(tabs, value=self.assistant_tab).classes('ai-assistant-tab-panels flex-grow'):
# Tab 1: Assistant (Chat)
with ui.tab_panel(self.assistant_tab).classes('flex flex-col w-full h-full'):
self._build_chat_tab(initial_message)
# Tab 2: Recent Conversations
with ui.tab_panel(conversations_tab).classes('flex flex-col h-full min-h-0'):
self._build_conversations_tab()
# Tab 3: Configuration
with ui.tab_panel(config_tab).classes('flex flex-col h-full'):
self._build_config_tab()
# Tab 4: Prompt Preview
with ui.tab_panel(preview_tab).classes('flex flex-col h-full'):
self._build_preview_tab()
def _build_chat_tab(self, initial_message: str):
"""Build the chat tab content."""
with ui.element('div').classes('flex flex-col w-full h-full'):
# Chat history (scrollable)
self.chat_history = ChatHistory(initial_message=initial_message)
# Chat input at bottom
self.input_el = chat_input(
placeholder='Ask your assistant...',
on_send=self._handle_send
)
async def _handle_send(self, message: str):
"""Handle message from user."""
self.add_user_message(message)
if self.on_send:
await self.on_send(message)
def _handle_mode_selector_change(self, e):
"""Handle mode change from dropdown selector (Assistant tab)."""
display_name = e.value
mode_config = self._stage_modes.get(display_name, {})
new_mode = mode_config.get("key", "standard")
# Only show notification if mode actually changed
if new_mode != self.output_mode:
self.output_mode = new_mode
# Sync to Configuration tab radio
if self.mode_radio:
self.mode_radio.set_value(display_name)
# Add system message to chat history
if self.chat_history:
self.chat_history.add_system_message(
f"Switched to {display_name} mode."
)
else:
self.output_mode = new_mode
# Update description and system prompt (same logic as config tab)
if self.prompt_description_label:
self.prompt_description_label.set_text(
mode_config.get("description", "")
)
# Handle custom mode
if self.system_prompt_input:
if new_mode == "custom":
self.system_prompt_input.props(remove='readonly')
else:
self.system_prompt_input.props(add='readonly')
# Update system prompt display
self.system_prompt_input.set_value(
mode_config.get("system_prompt", "")
)
def _build_conversations_tab(self):
"""Build the recent conversations tab content."""
with ui.element('div').classes('flex-grow overflow-y-auto min-h-0') as conv_list:
self.conversations_list = conv_list
self._render_conversations_list()
def _render_conversations_list(self):
"""Render the list of conversations."""
if self.conversations_list:
self.conversations_list.clear()
with self.conversations_list:
if not self.conversations:
ui.label('No previous conversations').classes('text-slate-500 text-sm p-4')
return
for conv in self.conversations:
conv_id = conv.get('id')
raw_title = conv.get('title') or 'Untitled conversation'
conv_stage = conv.get('stage', 'stage1')
stage_display = get_stage_display_name(conv_stage)
# Prepend stage name to title
title = f"{stage_display}: {raw_title}"
updated_at = conv.get('updated_at', '')
formatted_time = format_relative_timestamp(updated_at) if updated_at else ''
with ui.element('div').classes(
'flex items-center gap-3 py-3 px-2 hover:bg-slate-50 cursor-pointer border-b border-slate-100 group'
):
# Clickable area for loading conversation
with ui.element('div').classes('flex items-center gap-3 flex-grow min-w-0').on(
'click', lambda e, cid=conv_id, cs=conv_stage: self._handle_load_conversation(cid, cs)
):
# Chat icon
ui.icon('chat_bubble').classes('text-primary text-xl flex-shrink-0')
# Title with stage prefix
ui.label(title).classes('flex-grow text-sm font-medium text-slate-800 truncate')
# Timestamp and delete button
ui.label(formatted_time).classes('text-xs text-slate-500 whitespace-nowrap')
ui.button(
icon='delete_outline',
on_click=lambda e, cid=conv_id, t=raw_title: self._confirm_delete_conversation(cid, t)
).props('flat round dense size=sm').classes(
'opacity-0 group-hover:opacity-100 text-slate-400 hover:text-red-500 ml-1'
)
async def _handle_new_chat(self):
"""Handle starting a new chat."""
self.current_conversation_id = None
if self.chat_history:
self.chat_history.clear()
if self.on_new_chat:
self.on_new_chat()
# Switch to assistant tab
if self.tabs and self.assistant_tab:
self.tabs.set_value(self.assistant_tab)
async def _handle_load_conversation(self, conversation_id: str, conversation_stage: str = None):
"""Handle loading a conversation, navigating to correct stage if needed."""
# If conversation is from a different stage, navigate there first
if conversation_stage and conversation_stage != self.stage_context:
if self.on_navigate:
self.on_navigate(conversation_stage, conversation_id)
return
# Same stage - load normally
self.current_conversation_id = conversation_id
if self.on_load_conversation:
await self.on_load_conversation(conversation_id)
# Switch to assistant tab
if self.tabs and self.assistant_tab:
self.tabs.set_value(self.assistant_tab)
def _confirm_delete_conversation(self, conversation_id: str, title: str):
"""Show confirmation dialog before deleting a conversation."""
with ui.dialog() as dialog, ui.card().classes('p-4'):
ui.label(f'Delete "{title}"?').classes('text-lg font-medium mb-2')
ui.label('This will permanently delete this conversation and all its messages.').classes(
'text-sm text-slate-600 mb-4'
)
with ui.row().classes('gap-2 justify-end w-full'):
ui.button('Cancel', on_click=dialog.close).props('flat')
ui.button(
'Delete',
on_click=lambda: self._execute_delete(conversation_id, dialog)
).classes('bg-red-500 text-white')
dialog.open()
async def _execute_delete(self, conversation_id: str, dialog):
"""Execute the deletion after confirmation."""
dialog.close()
await self._handle_delete_conversation(conversation_id)
async def _handle_delete_conversation(self, conversation_id: str):
"""Handle deleting a conversation."""
if self.on_delete_conversation:
success = await self.on_delete_conversation(conversation_id)
if success:
# Remove from local list and re-render
self.conversations = [c for c in self.conversations if c.get('id') != conversation_id]
self._render_conversations_list()
# If we deleted the current conversation, start fresh
if self.current_conversation_id == conversation_id:
await self._handle_new_chat()
def update_conversations(self, conversations: list[dict]):
"""Update the conversations list."""
self.conversations = conversations
self._render_conversations_list()
def add_conversation_to_list(self, conversation: dict):
"""Add a new conversation to the top of the list."""
self.conversations.insert(0, conversation)
self._render_conversations_list()
def _build_config_tab(self):
"""Build the configuration tab content."""
# Get stage-specific configuration
stage_config = get_stage_config(self.stage_context)
modes = stage_config.get("modes", [])
default_mode = stage_config.get("default_mode", "standard")
# Build options list for radio
mode_options = [m["display_name"] for m in modes]
# Find default display name
default_display = next(
(m["display_name"] for m in modes if m["key"] == default_mode),
mode_options[0] if mode_options else "Standard"
)
# Store modes for lookup
self._stage_modes = {m["display_name"]: m for m in modes}
with ui.element('div').classes('px-4 pt-4 pb-2 h-full flex flex-col'):
# Output Mode Selection
with ui.element('div').classes('config-section'):
ui.label('Output Mode').classes('config-label')
self.mode_radio = ui.radio(
options=mode_options,
value=default_display,
on_change=self._handle_mode_change
).props('inline')
# Prompt Description (read-only text display)
with ui.element('div').classes('config-section'):
ui.label('Prompt Description').classes('config-label')
# Get initial description
initial_mode = self._stage_modes.get(default_display, {})
self.prompt_description_label = ui.label(
initial_mode.get("description", "Select an output mode.")
).classes('text-secondary text-sm')
# System Prompt - flex-grow to fill remaining space
with ui.element('div').classes('config-section flex-grow flex flex-col min-h-0'):
ui.label('System Prompt').classes('config-label')
self.system_prompt_input = ui.textarea(
placeholder='Enter system prompt...',
value=initial_mode.get("system_prompt", "")
).classes('config-textarea').props('readonly')
self._system_prompt_readonly = True
def _handle_mode_change(self, e):
"""Handle output mode radio button change (Configuration tab)."""
display_name = e.value
mode_config = self._stage_modes.get(display_name, {})
new_mode = mode_config.get("key", "standard")
# Only show notification if mode actually changed
if new_mode != self.output_mode:
self.output_mode = new_mode
# Sync to Assistant tab dropdown
if self.mode_selector:
self.mode_selector.set_value(display_name)
# Add system message to chat history
if self.chat_history:
self.chat_history.add_system_message(
f"Switched to {display_name} mode."
)
else:
self.output_mode = new_mode
# Update description
self.prompt_description_label.set_text(
mode_config.get("description", "")
)
# Handle custom mode
if new_mode == "custom":
self.system_prompt_input.props(remove='readonly')
else:
self.system_prompt_input.props(add='readonly')
# Update system prompt display
self.system_prompt_input.set_value(
mode_config.get("system_prompt", "")
)
def _build_preview_tab(self):
"""Build the prompt preview tab content."""
placeholder_prompt = """=== SYSTEM PROMPT ===
You are an educational assessment expert. Generate clear, measurable rubric criteria based on the provided prompt groups and task context.
=== USER CONTEXT ===
Project: [Project Title]
Task: [Task Name]
Prompt Groups:
1. [Group 1 Name]
- Description: [Description]
2. [Group 2 Name]
- Description: [Description]
=== INSTRUCTIONS ===
Generate a rubric with the following structure:
- Each prompt group becomes a criterion
- Include 4 performance levels (Exemplary, Proficient, Developing, Beginning)
- Provide clear, measurable descriptors for each level"""
with ui.element('div').classes('prompt-preview-container') as container:
self.prompt_preview_container = container
ui.label(placeholder_prompt).classes('whitespace-pre-wrap')
# Public API methods (same as ChatAssistant)
def add_ai_message(self, message: str):
"""Add an AI message to the chat history."""
self.hide_loading()
if self.chat_history:
self.chat_history.add_ai_message(message)
def add_user_message(self, message: str):
"""Add a user message to the chat history."""
if self.chat_history:
self.chat_history.add_user_message(message)
def show_loading(self):
"""Show a loading spinner in the chat."""
if self.chat_history and not self.current_spinner:
self.current_spinner = self.chat_history.add_spinner()
def hide_loading(self):
"""Hide the loading spinner."""
if self.current_spinner:
try:
self.current_spinner.delete()
except (ValueError, AttributeError):
# Spinner already removed or doesn't exist
pass
finally:
self.current_spinner = None