Feature/sotopia demo UI

sotopia/ui/streamlit_ui/pages/render_chat_websocket_simplified.py

@@ -5,10 +5,11 @@
 import aiohttp
 import requests
 import streamlit as st
+import time
 
 from sotopia.ui.streamlit_ui.utils import get_abstract
-
-chat_container = st.empty()
+from sotopia.ui.streamlit_ui.rendering.render_episode import rendering_episode_full
+from sotopia.database import EpisodeLog
 
 
 def compose_agent_names(agent_dict: dict[Any, Any]) -> str:
@@ -61,14 +62,16 @@ def initialize_session_state() -> None:
         print("Session state initialized")
 
 
-def parse_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+chat_history_container = st.empty()
+
+
+def parse_messages(messages: list[str]) -> list[dict[str, Any]]:
     chat_messages = []
-    messages = messages[1:]
+    messages = [message for message in messages[1:] if message != ""]
     evaluation_available = len(messages) >= 2 and "Agent 1 comments" in messages[-2]
     evaluations = [] if not evaluation_available else messages[-2:]
     conversations = messages if not evaluation_available else messages[:-2]
 
-    # print(messages)
     chat_messages = [
         {"role": "human", "content": message} for message in conversations
     ] + [
@@ -124,27 +127,25 @@ async def run_simulation():
                         break
 
                     if message["data"]["type"] == "messages":
-                        st.session_state.messages = parse_messages(
-                            message["data"]["messages"]
-                        )
+                        epilog = EpisodeLog(**message["data"]["messages"])
+                        st.session_state.messages = epilog
+                        with chat_history_container.container():
+                            rendering_episode_full(epilog)
+
                 elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
                     break
 
                 if st.session_state.get("stop_sim", False):
                     await ws.send_str(json.dumps({"type": "FINISH_SIM", "data": ""}))
 
-                with chat_container.container():
-                    for msg in st.session_state.messages:
-                        st.chat_message(msg["role"]).write(msg["content"])
-
 
 def start_callback() -> None:
     if st.session_state.agent_choice_1 == st.session_state.agent_choice_2:
         st.error("Please select different agents")
     else:
         st.session_state.active = True
         st.session_state.stop_sim = False
-        st.session_state.messages = []
+        st.session_state.messages = None
         try:
             asyncio.run(run_simulation())
         except Exception as e:
@@ -155,12 +156,6 @@ def start_callback() -> None:
 
 def stop_callback() -> None:
     st.session_state.stop_sim = True
-    st.session_state.websocket_manager.send_message(
-        {
-            "type": "FINISH_SIM",
-            "data": "",
-        }
-    )
 
 
 def is_active() -> bool:
@@ -240,9 +235,22 @@ def chat_demo() -> None:
                 on_click=stop_callback,
             )
 
-        with chat_container.container():
-            for msg in st.session_state.messages:
-                st.chat_message(msg["role"]).write(msg["content"])
+        # if is_active():
+        #     try:
+        #         asyncio.run(run_simulation())
+        #     except Exception as e:
+        #         st.error(f"Error running simulation: {e}")
+        #     finally:
+        #         st.session_state.active = False
+
+        while is_active() and st.session_state.messages:
+            with chat_history_container.container():
+                rendering_episode_full(st.session_state.messages)
+            time.sleep(1)
+
+        if st.session_state.messages:
+            with chat_history_container.container():
+                rendering_episode_full(st.session_state.messages)
 
 
 chat_demo()

sotopia/ui/streamlit_ui/rendering/render_episode.py

@@ -92,6 +92,92 @@ def rendering_episode(episode: EpisodeLog) -> None:
             )
 
 
+def rendering_episode_full(episode: EpisodeLog) -> None:
+    print("Episode: ", episode)
+    agents = [AgentProfile.get(agent) for agent in episode.agents]
+    agent_names = [get_full_name(agent) for agent in agents]
+    environment = EnvironmentProfile.get(episode.environment)
+
+    avatar_mapping = {
+        agent_names[0]: "🧔🏻",
+        agent_names[1]: "🧑",
+    }
+
+    messages = render_for_humans(episode)
+
+    background_messages = [
+        message for message in messages if message["role"] == "Background Info"
+    ]
+    evaluation_messages = [
+        message for message in messages if message["type"] == "comment"
+    ]
+    conversation_messages = [
+        message
+        for message in messages
+        if message not in background_messages and message not in evaluation_messages
+    ]
+
+    assert (
+        len(background_messages) == 2
+    ), f"Need 2 background messages, but got {len(background_messages)}"
+
+    print(f"\n\ENVIRONMENT {environment}")
+
+    rendering_episode(episode)
+
+    st.markdown("---")
+
+    st.subheader("Conversation & Evaluation")
+    with st.expander("Conversation", expanded=True):
+        for index, message in enumerate(conversation_messages):
+            role = role_mapping.get(message["role"], message["role"])
+            content = message["content"]
+
+            if role == "obs" or message.get("type") == "action":
+                try:
+                    content = json.loads(content)
+                except Exception as e:
+                    print(e)
+
+            with st.chat_message(
+                role, avatar=avatar_mapping.get(message["role"], None)
+            ):
+                if isinstance(content, dict):
+                    st.json(content)
+                elif role == "info":
+                    st.markdown(
+                        f"""
+                        <div style="background-color: lightblue; padding: 10px; border-radius: 5px;">
+                            {content}
+                        </div>
+                        """,
+                        unsafe_allow_html=True,
+                    )
+                else:
+                    if role == agent_names[1]:
+                        # add background grey color
+                        st.write(f"**{role}**")
+                        st.markdown(
+                            f"""
+                            <div style="background-color: lightgrey; padding: 10px; border-radius: 5px;">
+                                {content}
+                            </div>
+                            """,
+                            unsafe_allow_html=True,
+                        )
+                    else:
+                        st.write(f"**{role}**")
+                        st.markdown(
+                            content.replace("\n", "<br />"), unsafe_allow_html=True
+                        )
+
+    with st.expander("Evaluation Results", expanded=True):
+        for message in evaluation_messages:
+            st.markdown(
+                message["content"].replace("\n", "<br />"), unsafe_allow_html=True
+            )
+
+
 def rendering_episodes() -> None:
     local_css("./././css/style.css")
 

sotopia/ui/websocket_utils.py

@@ -158,11 +158,11 @@ async def arun(self) -> AsyncGenerator[dict[str, Any], None]:
 
         async for messages in await generator:  # type: ignore
             reasoning, rewards = "", [0.0, 0.0]
-            eval_available = False
+            # eval_available = False
             if messages[-1][0][0] == "Evaluation":
                 reasoning = messages[-1][0][2].to_natural_language()
                 rewards = eval(messages[-2][0][2].to_natural_language())
-                eval_available = True
+                # eval_available = True
 
             epilog = EpisodeLog(
                 environment=self.env.profile.pk,
@@ -177,11 +177,11 @@ async def arun(self) -> AsyncGenerator[dict[str, Any], None]:
                 rewards=rewards,
                 rewards_prompt="",
             )
-            agent_profiles, parsed_messages = epilog.render_for_humans()
-            if not eval_available:
-                parsed_messages = parsed_messages[:-2]
+            # agent_profiles, parsed_messages = epilog.render_for_humans()
+            # if not eval_available:
+            #     parsed_messages = parsed_messages[:-2]
 
             yield {
                 "type": "messages",
-                "messages": parsed_messages,
+                "messages": epilog.dict(),
             }