Fix issues with local image as st.chat_message avatar

e2e/playwright/st_chat_message.py

@@ -0,0 +1,80 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pandas as pd
+from PIL import Image
+
+import streamlit as st
+
+np.random.seed(0)
+
+
+# Generate a random dataframe
+df = pd.DataFrame(
+    np.random.randn(5, 5),
+    columns=("col_%d" % i for i in range(5)),
+)
+
+
+with st.chat_message("user"):
+    st.write("Hello…")
+
+with st.chat_message("assistant"):
+    st.write(
+        """
+Hello, here is a code snippet:
+
+```python
+import streamlit as st
+with st.chat_message("assistant"):
+     st.write("Hello, here is a code snippet...")
+```
+"""
+    )
+
+with st.chat_message("user", avatar="🧑"):
+    st.write(
+        """
+Lorem ipsum dolor sit amet, consectetur adipiscing elit. Mauris tristique est
+at tincidunt pul vinar. Nam pulvinar neque sapien, eu pellentesque metus pellentesque
+at. Ut et dui molestie, iaculis magna sed.
+"""
+    )
+
+with st.chat_message("dog", avatar="https://static.streamlit.io/examples/dog.jpg"):
+    st.write("Woof woof! I'm a dog and I like charts:")
+    st.line_chart(df, use_container_width=True)
+
+cat = st.chat_message("cat", avatar="https://static.streamlit.io/examples/cat.jpg")
+cat.write("I'm a cat and I like this dataset:")
+cat.dataframe(df, use_container_width=True)
+cat.text_input("What's your name?")
+
+
+with st.chat_message("Bot"):
+    with st.expander("See more", expanded=True):
+        st.write("Lorem ipsum dolor sit amet")
+
+st.chat_message("human")
+
+
+image1 = Image.new("RGB", (10, 10), "red")
+st.chat_message("user", avatar=image1).write("Red local image")
+
+image2 = Image.new("RGB", (10, 10), "blue")
+st.chat_message("assistant", avatar=image2).write("Blue local image")
+st.chat_message("assistant", avatar=image2).write(
+    "Another message with the same blue avatar."
+)

e2e/playwright/st_chat_message_test.py

@@ -0,0 +1,31 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from playwright.sync_api import Page, expect
+
+from conftest import ImageCompareFunction
+
+
+def test_renders_chat_messages_correctly_1(
+    themed_app: Page, assert_snapshot: ImageCompareFunction
+):
+    """Test if the chat messages render correctly"""
+    # Wait a bit more to allow all images to load:
+    chat_message_elements = themed_app.locator(".stChatMessage")
+    expect(chat_message_elements).to_have_count(10)
+    for i, element in enumerate(chat_message_elements.all()):
+        element.scroll_into_view_if_needed()
+        # Wait a bit more to allow the avatar images to load:
+        themed_app.wait_for_timeout(100)
+        assert_snapshot(element, name=f"chat_message-{i}")

frontend/lib/src/components/core/Block/Block.tsx

@@ -112,6 +112,7 @@ const BlockNodeRenderer = (props: BlockPropsWithWidth): ReactElement => {
     return (
       <ChatMessage
         element={node.deltaBlock.chatMessage as BlockProto.ChatMessage}
+        endpoints={props.endpoints}
       >
         {child}
       </ChatMessage>

frontend/lib/src/components/elements/ChatMessage/ChatMessage.test.tsx

@@ -18,6 +18,7 @@ import React from "react"
 import "@testing-library/jest-dom"
 import { screen } from "@testing-library/react"
 
+import { mockEndpoints } from "@streamlit/lib/src/mocks/mocks"
 import { render } from "@streamlit/lib/src/test_util"
 import { Block as BlockProto } from "@streamlit/lib/src/proto"
 
@@ -32,6 +33,9 @@ const getProps = (
     avatar: "user",
     ...elementProps,
   }),
+  endpoints: mockEndpoints({
+    buildMediaURL: jest.fn().mockImplementation(url => url),
+  }),
 })
 
 describe("ChatMessage", () => {

frontend/lib/src/components/elements/ChatMessage/ChatMessage.tsx

@@ -21,6 +21,7 @@ import { Face, SmartToy } from "@emotion-icons/material-outlined"
 import { Block as BlockProto } from "@streamlit/lib/src/proto"
 import Icon from "@streamlit/lib/src/components/shared/Icon"
 import { EmotionTheme } from "@streamlit/lib/src/theme"
+import { StreamlitEndpoints } from "@streamlit/lib/src/StreamlitEndpoints"
 
 import {
   StyledChatMessageContainer,
@@ -34,16 +35,22 @@ interface ChatMessageAvatarProps {
   name: string
   avatar?: string
   avatarType?: BlockProto.ChatMessage.AvatarType
+  endpoints: StreamlitEndpoints
 }
 
 function ChatMessageAvatar(props: ChatMessageAvatarProps): ReactElement {
-  const { avatar, avatarType, name } = props
+  const { avatar, avatarType, name, endpoints } = props
   const theme: EmotionTheme = useTheme()
 
   if (avatar) {
     switch (avatarType) {
       case BlockProto.ChatMessage.AvatarType.IMAGE:
-        return <StyledAvatarImage src={avatar} alt={`${name} avatar`} />
+        return (
+          <StyledAvatarImage
+            src={endpoints.buildMediaURL(avatar)}
+            alt={`${name} avatar`}
+          />
+        )
       case BlockProto.ChatMessage.AvatarType.EMOJI:
         return <StyledAvatarBackground>{avatar}</StyledAvatarBackground>
       case BlockProto.ChatMessage.AvatarType.ICON:
@@ -78,10 +85,12 @@ function ChatMessageAvatar(props: ChatMessageAvatarProps): ReactElement {
 }
 
 export interface ChatMessageProps {
+  endpoints: StreamlitEndpoints
   element: BlockProto.ChatMessage
 }
 
 const ChatMessage: React.FC<ChatMessageProps> = ({
+  endpoints,
   element,
   children,
 }): ReactElement => {
@@ -93,7 +102,12 @@ const ChatMessage: React.FC<ChatMessageProps> = ({
       data-testid="stChatMessage"
       background={["user", "human"].includes(name.toLowerCase())}
     >
-      <ChatMessageAvatar name={name} avatar={avatar} avatarType={avatarType} />
+      <ChatMessageAvatar
+        name={name}
+        avatar={avatar}
+        avatarType={avatarType}
+        endpoints={endpoints}
+      />
       <StyledMessageContent
         data-testid="stChatMessageContent"
         aria-label={`Chat message from ${name}`}

lib/streamlit/elements/widgets/chat.py

@@ -51,14 +51,17 @@ class PresetNames(str, Enum):
 
 
 def _process_avatar_input(
-    avatar: str | AtomicImage | None = None,
+    avatar: str | AtomicImage | None, delta_path: str
 ) -> Tuple[BlockProto.ChatMessage.AvatarType.ValueType, str]:
     """Detects the avatar type and prepares the avatar data for the frontend.
 
     Parameters
     ----------
     avatar :
         The avatar that was provided by the user.
+    delta_path : str
+        The delta path is used as media ID when a local image is served via the media
+        file manager.
 
     Returns
     -------
@@ -89,7 +92,7 @@ def _process_avatar_input(
                 clamp=False,
                 channels="RGB",
                 output_format="auto",
-                image_id="",
+                image_id=delta_path,
             )
         except Exception as ex:
             raise StreamlitAPIException(
@@ -194,7 +197,9 @@ def chat_message(
         ):
             # For selected labels, we are mapping the label to an avatar
             avatar = name.lower()
-        avatar_type, converted_avatar = _process_avatar_input(avatar)
+        avatar_type, converted_avatar = _process_avatar_input(
+            avatar, self.dg._get_delta_path_str()
+        )
 
         message_container_proto = BlockProto.ChatMessage()
         message_container_proto.name = name