Add simple chat widget function.

Also adds an example usage of the chat widget.

Functionality:

- Provide title
- Provide a bot name
- Can stream chat response
- When generating a response, new prompts cannot
- be sent
- Relatively responsive to device width

Minor issues:

- Chat window does not scroll along with new text being added.
- Text input does not clear when pressing entering a prompt.

Author: richard-to

mesop/examples/__init__.py

@@ -1,4 +1,5 @@
 from mesop.examples import buttons as buttons
+from mesop.examples import chat as chat
 from mesop.examples import columns as columns
 from mesop.examples import composite as composite
 from mesop.examples import docs as docs

mesop/examples/chat.py

@@ -0,0 +1,40 @@
+import random
+import time
+
+import mesop as me
+import mesop.labs as mel
+
+
+def lorem_ipsum_chat(input):
+  lines = [
+    (
+      "Lorem ipsum dolor sit amet, consectetur adipiscing elit, "
+      "sed do eiusmod tempor incididunt ut labore et dolore magna aliqua."
+    ),
+    "Laoreet sit amet cursus sit amet dictum sit amet.",
+    "At lectus urna duis convallis.",
+    "A pellentesque sit amet porttitor eget.",
+    "Mauris nunc congue nisi vitae suscipit tellus mauris a diam.",
+    "Aliquet lectus proin nibh nisl condimentum id.",
+    "Integer malesuada nunc vel risus commodo viverra maecenas accumsan.",
+    "Tempor id eu nisl nunc mi.",
+    "Id consectetur purus ut faucibus pulvinar.",
+    "Mauris pharetra et ultrices neque ornare.",
+    "Facilisis magna etiam tempor orci.",
+    "Mauris pharetra et ultrices neque.",
+    "Sit amet facilisis magna etiam tempor orci.",
+    "Amet consectetur adipiscing elit pellentesque habitant morbi tristique.",
+    "Egestas erat imperdiet sed euismod.",
+    "Tincidunt praesent semper feugiat nibh sed pulvinar proin gravida.",
+    "Habitant morbi tristique senectus et netus et malesuada.",
+  ]
+  for line in random.sample(lines, random.randint(3, len(lines) - 1)):
+    time.sleep(0.5)
+    yield line + " "
+
+
+@me.page(path="/")
+def page():
+  mel.chat(
+    lorem_ipsum_chat, title="Lorem Ipsum Chat", bot_username="Lorem Ipsum Bot"
+  )

mesop/labs/__init__.py

@@ -1 +1,2 @@
+from mesop.labs.chat import chat as chat
 from mesop.labs.io import text_io as text_io

mesop/labs/chat.py

@@ -0,0 +1,235 @@
+import time
+from dataclasses import dataclass
+from typing import Callable, Generator, Literal
+
+import mesop as me
+
+_ROLE_USER = "user"
+_ROLE_ASSISTANT = "assistant"
+
+_BOT_USERNAME_DEFAULT = "mesop-bot"
+
+_COLOR_BACKGROUND = "#f0f4f8"
+_COLOR_CHAT_BUBBLE_YOU = "#f2f2f2"
+_COLOR_CHAT_BUBBLE_BOT = "#ebf3ff"
+
+_DEFAULT_PADDING = me.Padding(top=20, left=20, right=20, bottom=20)
+_DEFAULT_BORDER_SIDE = me.BorderSide(
+  width="1px", style="solid", color="#ececec"
+)
+
+_LABEL_BUTTON = "Send prompt"
+_LABEL_BUTTON_IN_PROGRESS = "Processing prompt..."
+_LABEL_INPUT = "Enter your prompt"
+
+_STYLE_APP_CONTAINER = me.Style(
+  background=_COLOR_BACKGROUND,
+  display="grid",
+  height="100vh",
+  grid_template_columns="repeat(1, 1fr)",
+)
+_STYLE_TITLE = me.Style(padding=me.Padding(left=10))
+_STYLE_CHAT_BOX = me.Style(
+  height="100%",
+  overflow_y="scroll",
+  padding=_DEFAULT_PADDING,
+  margin=me.Margin(bottom="20px"),
+  border_radius="10px",
+  border=me.Border(
+    left=_DEFAULT_BORDER_SIDE,
+    right=_DEFAULT_BORDER_SIDE,
+    top=_DEFAULT_BORDER_SIDE,
+    bottom=_DEFAULT_BORDER_SIDE,
+  ),
+)
+_STYLE_CHAT_INPUT = me.Style(width="100%")
+_STYLE_CHAT_INPUT_BOX = me.Style(padding=me.Padding(top="30px"))
+_STYLE_CHAT_BUBBLE_NAME = me.Style(
+  font_weight="bold",
+  font_size="12px",
+  padding=me.Padding(left="15px", right="15px", bottom="5px"),
+)
+_STYLE_CHAT_BUBBLE_PLAINTEXT = me.Style(
+  margin=me.Margin(top="15px", bottom="15px")
+)
+
+
+def _make_style_chat_ui_container(has_title: bool) -> me.Style:
+  """Generates styles for chat UI container depending on if there is a title or not.
+
+  Args:
+    has_title: Whether the Chat UI is display a title or not.
+  """
+  return me.Style(
+    display="grid",
+    grid_template_columns="repeat(1, 1fr)",
+    grid_template_rows="1fr 14fr 1fr" if has_title else "5fr 1fr",
+    margin=me.Margin(top=0, bottom=0, left="auto", right="auto"),
+    width="min(1024px, 100%)",
+    height="100vh",
+    background="#fff",
+    box_shadow=(
+      "0 3px 1px -2px #0003, 0 2px 2px #00000024, 0 1px 5px #0000001f"
+    ),
+    padding=_DEFAULT_PADDING,
+  )
+
+
+def _make_style_chat_bubble_wrapper(
+  role: Literal["user", "assistant"]
+) -> me.Style:
+  """Generates styles for chat bubble position.
+
+  Args:
+    role: Chat bubble alignment depends on the role
+  """
+  align_items = "end" if role == _ROLE_USER else "start"
+  return me.Style(
+    display="flex",
+    flex_direction="column",
+    align_items=align_items,
+  )
+
+
+def _make_chat_bubble_style(role: Literal["user", "assistant"]) -> me.Style:
+  """Generates styles for chat bubble.
+
+  Args:
+    role: Chat bubble background color depends on the role
+  """
+  background = (
+    _COLOR_CHAT_BUBBLE_YOU if role == _ROLE_USER else _COLOR_CHAT_BUBBLE_BOT
+  )
+  return me.Style(
+    width="80%",
+    font_size="13px",
+    background=background,
+    border_radius="15px",
+    padding=me.Padding(
+      top="0px",
+      right="15px",
+      left="15px",
+      bottom="3px",
+    ),
+    margin=me.Margin(bottom="10px"),
+    border=me.Border(
+      left=_DEFAULT_BORDER_SIDE,
+      right=_DEFAULT_BORDER_SIDE,
+      top=_DEFAULT_BORDER_SIDE,
+      bottom=_DEFAULT_BORDER_SIDE,
+    ),
+  )
+
+
+@dataclass
+class Message:
+  """Chat message metadata."""
+
+  role: Literal["assistant", "user"] = "user"
+  username: str = ""
+  content: str = ""
+
+
+@me.stateclass
+class State:
+  input: str
+  output: list[Message]
+  in_progress: bool = False
+
+
+def on_input_update(State):
+  """Generic on input handler that saves input to State using the given key.
+
+  This helper only works if you have one state instance. If use multiple state classes
+  with this helper, then only the last event handler will be stored. For more info, see
+  https://google.github.io/mesop/guides/troubleshooting/#avoid-using-closure-variables-in-event-handler.
+  """
+
+  def _on_update(
+    e: (
+      me.InputEvent
+      | me.SelectSelectionChangeEvent
+      | me.RadioChangeEvent
+      | me.CheckboxChangeEvent
+    ),
+  ):
+    state = me.state(State)
+    setattr(state, e.key, e.value)
+
+  return _on_update
+
+
+def chat(
+  transform: Callable[[str], Generator[str, None, None] | str],
+  *,
+  title: str | None = None,
+  bot_username: str = _BOT_USERNAME_DEFAULT,
+):
+  """Creates a simple chat UI which takes in a prompt and returns a response to the
+  prompt.
+
+  This function creates event handlers for text input and output operations
+  using the provided function `transform` to process the input and generate the output.
+
+  Args:
+    transform: Function that takes in a prompt and returns a response to the prompt.
+    title: Headline text to display at the top of the UI.
+    bot_username: Name of your bot / assistant.
+  """
+  state = me.state(State)
+
+  def on_click(e: me.ClickEvent):
+    state = me.state(State)
+    if state.in_progress or not state.input:
+      return
+
+    output = state.output
+    if output is None:
+      output = []
+    output.append(Message(role=_ROLE_USER, content=state.input))
+    state.in_progress = True
+    yield
+
+    start_time = time.time()
+    output_message = transform(state.input)
+    assistant_message = Message(role=_ROLE_ASSISTANT, username=bot_username)
+    output.append(assistant_message)
+    state.output = output
+    for content in output_message:
+      assistant_message.content += content
+      if (time.time() - start_time) >= 0.5:
+        start_time = time.time()
+        yield
+    state.in_progress = False
+    yield
+
+  with me.box(style=_STYLE_APP_CONTAINER):
+    with me.box(style=_make_style_chat_ui_container(bool(title))):
+      if title:
+        me.text(title, type="headline-5", style=_STYLE_TITLE)
+      with me.box(style=_STYLE_CHAT_BOX):
+        for msg in state.output:
+          with me.box(style=_make_style_chat_bubble_wrapper(msg.role)):
+            if msg.role == _ROLE_ASSISTANT:
+              me.text(msg.username, style=_STYLE_CHAT_BUBBLE_NAME)
+            with me.box(style=_make_chat_bubble_style(msg.role)):
+              if msg.role == _ROLE_USER:
+                me.text(msg.content, style=_STYLE_CHAT_BUBBLE_PLAINTEXT)
+              else:
+                me.markdown(msg.content)
+
+      with me.box(style=_STYLE_CHAT_INPUT_BOX):
+        me.input(
+          label=_LABEL_INPUT,
+          key="input",
+          on_input=on_input_update(State),
+          style=_STYLE_CHAT_INPUT,
+        )
+        with me.box():
+          me.button(
+            _LABEL_BUTTON_IN_PROGRESS if state.in_progress else _LABEL_BUTTON,
+            color="primary",
+            type="flat",
+            disabled=state.in_progress,
+            on_click=on_click,
+          )