Feature/sotopia demo UI

docs/pages/contribution/contribution.md

@@ -133,7 +133,7 @@ Please refer to [Dev Containers](https://containers.dev/supporting#editors) to s
 
 You can also set up the development environment without Dev Containers. There are three things you will need to set up manually:
 
-- Python and uv: Please start from an environment supporting Python 3.10+ and install uv using `pip install uv; uv sync --all-extra`.
+- Python and uv: Please start from an environment supporting Python 3.10+ and install uv using `pip install uv; uv sync --all-extras`. (Note that this will install all the extra dependencies)
 - Redis: Please refer to introduction page for the set up of Redis.
 - Local LLM (optional): If you don't have access to model endpoints (e.g. OpenAI, Anthropic or others), you can use a local model. You can use Ollama, Llama.cpp,  vLLM or many others which support OpenAI compatible endpoints.
 

pyproject.toml

@@ -40,6 +40,8 @@ examples = ["transformers", "datasets", "scipy", "torch", "pandas"]
 api = [
     "fastapi[standard]",
     "uvicorn",
+    "streamlit",
+    "websockets"
 ]
 test = ["pytest", "pytest-cov", "pytest-asyncio"]
 

sotopia/ui/README.md

@@ -2,6 +2,14 @@
 > [!CAUTION]
 > Work in progress: the API endpoints are being implemented. And will be released in the future major version.
 
+## Streamlit UI
+To run the Streamlit UI, run the following command:
+```bash
+cd sotopia/ui/streamlit_ui
+uv run streamlit run app.py
+```
+
+
 ## FastAPI Server
 
 The API server is a FastAPI application that is used to connect the Sotopia UI to the Sotopia backend.

sotopia/ui/streamlit_ui/.streamlit/config.toml

@@ -0,0 +1,11 @@
+[theme]
+base="light"
+primaryColor = "#26184e"
+backgroundColor="#FFFFFF"
+secondaryBackgroundColor="#F0F2F6"
+textColor="#31333F"
+font="sans serif"
+
+
+[server]
+enableStaticServing = true

sotopia/ui/streamlit_ui/__init__.py

@@ -0,0 +1,7 @@
+from .utils import reset_database
+from .rendering.rendering_utils import (
+    _agent_profile_to_friendabove_self,
+    render_for_humans,
+)
+
+__all__ = ["reset_database", "_agent_profile_to_friendabove_self", "render_for_humans"]

sotopia/ui/streamlit_ui/app.py

@@ -0,0 +1,95 @@
+import os
+
+import streamlit as st
+
+from sotopia.ui.streamlit_ui.utils import reset_database
+
+PORT = 8800
+st.session_state.API_BASE = f"http://localhost:{PORT}"
+st.session_state.WS_BASE = f"ws://localhost:{PORT}"
+
+
+def update_database_callback() -> None:
+    new_database_url = st.session_state.new_database_url
+    updated_url = (
+        new_database_url if new_database_url != "" else st.session_state.DEFAULT_DB_URL
+    )
+    try:
+        reset_database(updated_url)
+    except Exception as e:
+        st.error(f"Error occurred while updating database: {e}, please try again.")
+
+    st.session_state.current_database_url = updated_url
+    print("Updated DB URL: ", st.session_state.current_database_url)
+
+
+# Page Configuration
+st.set_page_config(page_title="SocialStream_Demo", page_icon="🧊", layout="wide")
+
+display_intro = st.Page(
+    "./pages/intro.py", title="Introduction", icon=":material/home:"
+)
+display_episodes = st.Page(
+    "./pages/display_episodes.py", title="Episode", icon=":material/photo_library:"
+)
+display_scenarios = st.Page(
+    "./pages/display_scenarios.py",
+    title="Scenarios",
+    icon=":material/insert_drive_file:",
+)
+display_characters = st.Page(
+    "./pages/display_characters.py", title="Characters", icon=":material/people:"
+)
+
+display_chat = st.Page(
+    "./pages/render_chat_websocket_simplified.py",
+    title="Simulation",
+    icon=":material/add:",
+)
+
+add_characters = st.Page(
+    "./pages/add_characters.py", title="Add Characters", icon=":material/add:"
+)
+
+add_scenarios = st.Page(
+    "./pages/add_scenarios.py", title="Add Scenarios", icon=":material/add:"
+)
+
+pg = st.navigation(
+    [
+        display_intro,
+        display_scenarios,
+        display_episodes,
+        display_characters,
+        display_chat,
+        add_characters,
+        add_scenarios,
+    ]
+)
+
+# Reset active agent when switching modes across pages
+if "mode" not in st.session_state or pg.title != st.session_state.get("mode", None):
+    if "active" in st.session_state:
+        del st.session_state["active"]
+        # print("Active agent reset.")
+
+    st.session_state.mode = pg.title
+
+
+# DB URL Configuration
+if "DEFAULT_DB_URL" not in st.session_state:
+    st.session_state.DEFAULT_DB_URL = os.environ.get("REDIS_OM_URL", "")
+    st.session_state.current_database_url = st.session_state.DEFAULT_DB_URL
+    print("Default DB URL: ", st.session_state.DEFAULT_DB_URL)
+
+# impl 2: popup update URL
+with st.sidebar.popover("(Optional) Enter Database URL"):
+    new_database_url = st.text_input(
+        "URL: (starting in redis://)",
+        value="",
+        on_change=update_database_callback,
+        key="new_database_url",
+    )
+
+
+pg.run()

sotopia/ui/streamlit_ui/css/style.css

@@ -0,0 +1,28 @@
+.truncate{
+    max-height: 2em;
+    text-overflow: ellipsis;
+    margin-bottom: 12px;
+    cursor: pointer;
+    word-break: break-all;
+    overflow:hidden;
+    white-space: nowrap;
+    transition: max-height 1s, white-space 1s;
+}
+.truncate:hover{
+    overflow: auto;
+    overflow-y: scroll;
+    white-space: normal;
+    max-height: 200px;
+}
+/*
+[data-testid="stAppViewContainer"] {
+    background-image: url("https://images.unsplash.com/photo-1729096532452-50549f2bb63c");
+    background-size: 40%;
+    background-position: top right;
+    background-repeat: no-repeat;
+    background-attachment: local;
+} */
+
+[data-testid="stHeader"] {
+    background: rgba(0,0,0,0);
+}

sotopia/ui/streamlit_ui/pages/add_characters.py

@@ -0,0 +1,86 @@
+"""
+Definition
+@app.post("/agents/", response_model=str)
+async def create_agent(agent: AgentProfileWrapper) -> str:
+    agent_profile = AgentProfile(**agent.model_dump())
+    agent_profile.save()
+    pk = agent_profile.pk
+    assert pk is not None
+    return pk
+
+class AgentProfileWrapper(BaseModel):
+    Wrapper for AgentProfile to avoid pydantic v2 issues
+
+    first_name: str
+    last_name: str
+    age: int = 0
+    occupation: str = ""
+    gender: str = ""
+    gender_pronoun: str = ""
+    public_info: str = ""
+    big_five: str = ""
+    moral_values: list[str] = []
+    schwartz_personal_values: list[str] = []
+    personality_and_values: str = ""
+    decision_making_style: str = ""
+    secret: str = ""
+    model_id: str = ""
+    mbti: str = ""
+    tag: str = ""
+"""
+
+import streamlit as st
+import requests
+
+from sotopia.ui.fastapi_server import AgentProfileWrapper
+# add fields for agent profiles
+
+
+def local_css(file_name: str) -> None:
+    with open(file_name) as f:
+        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
+
+
+def rendering_character_form() -> None:
+    local_css("././css/style.css")
+    st.markdown("<h1>Character Creation</h1>", unsafe_allow_html=True)
+    st.write("Fill in the fields below to create a new character:")
+
+    first_name = st.text_input("First Name")
+    last_name = st.text_input("Last Name")
+    age = st.number_input("Age", min_value=0)
+    occupation = st.text_input("Occupation")
+    gender = st.text_input("Gender")
+    gender_pronoun = st.text_input("Gender Pronoun")
+    public_info = st.text_area("Public Info")
+
+    if st.button("Create Character"):
+        agent_profile = AgentProfileWrapper(
+            first_name=first_name,
+            last_name=last_name,
+            age=age,
+            occupation=occupation,
+            gender=gender,
+            gender_pronoun=gender_pronoun,
+            public_info=public_info,
+        )
+        print(agent_profile)
+
+        response = requests.post(
+            f"{st.session_state.API_BASE}/agents/",
+            json=agent_profile.model_dump(),
+        )
+
+        if response.status_code != 200:
+            st.error("Failed to create character. Error message: " + response.text)
+
+        else:
+            agent_id = response.json()
+            st.success("Character created successfully! ID: " + agent_id)
+            retrieved_agent = requests.get(
+                f"{st.session_state.API_BASE}/agents/id/{agent_id}"
+            )
+            st.write(retrieved_agent.json())
+
+
+rendering_character_form()

sotopia/ui/streamlit_ui/pages/add_scenarios.py

@@ -0,0 +1,123 @@
+"""
+Definition
+@app.post("/scenarios/", response_model=str)
+async def create_scenario(scenario: EnvironmentProfileWrapper) -> str:
+    scenario_profile = EnvironmentProfile(**scenario.model_dump())
+    scenario_profile.save()
+    pk = scenario_profile.pk
+    assert pk is not None
+    return pk
+
+class EnvironmentProfileWrapper(BaseModel):
+    Wrapper for EnvironmentProfile to avoid pydantic v2 issues
+
+    codename: str
+    source: str = ""
+    scenario: str = ""
+    agent_goals: list[str] = []
+    relationship: Literal[0, 1, 2, 3, 4, 5] = 0
+    age_constraint: str | None = None
+    occupation_constraint: str | None = None
+    agent_constraint: list[list[str]] | None = None
+    tag: str = ""
+
+class RelationshipType(IntEnum):
+    stranger = 0
+    know_by_name = 1
+    acquaintance = 2
+    friend = 3
+    romantic_relationship = 4
+    family_member = 5
+
+The age constraint of the environment, a list of tuples, each tuple is a range of age, e.g., '[(18, 25), (30, 40)]'
+means the environment is only available to agent one between 18 and 25, and agent two between 30 and 40
+"""
+
+import streamlit as st
+from sotopia.ui.fastapi_server import EnvironmentProfileWrapper
+from sotopia.database import RelationshipType
+import requests
+
+
+def local_css(file_name: str) -> None:
+    with open(file_name) as f:
+        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
+
+
+def rendering_scenario_form() -> None:
+    local_css("././css/style.css")
+    st.markdown("<h1>Scenario Creation</h1>", unsafe_allow_html=True)
+
+    codename = st.text_input("Codename")
+    source = st.text_input("Source")
+    scenario = st.text_area("Scenario")
+    # present relationship type with the descriptions, only accept one choice, then map to the enum
+    relationship_mapping = {
+        "Stranger": RelationshipType.stranger,
+        "Know by Name": RelationshipType.know_by_name,
+        "Acquaintance": RelationshipType.acquaintance,
+        "Friend": RelationshipType.friend,
+        "Romantic Relationship": RelationshipType.romantic_relationship,
+        "Family Member": RelationshipType.family_member,
+    }
+
+    selected_relationship = st.selectbox(
+        "Relationship",
+        list(relationship_mapping.keys()),
+    )
+    relationship = relationship_mapping[selected_relationship]
+
+    # first choose whether to use age constraint and then choose age range
+    use_age_constraint = st.checkbox("Use Age Constraint")
+    if use_age_constraint:
+        min_age = st.number_input("Min Age", min_value=0, max_value=100)
+        max_age = st.number_input("Max Age", min_value=0, max_value=100)
+        age_constraint = f"[({min_age}, {max_age})]"
+        if min_age > max_age:
+            st.error("Min age cannot be greater than max age")
+    else:
+        age_constraint = None
+
+    # first choose whether to use occupation constraint and then choose occupation
+    use_occupation_constraint = st.checkbox(
+        "Use Occupation Constraint, use comma to separate multiple occupations"
+    )
+    if use_occupation_constraint:
+        occupation = st.text_input("Occupation")
+        occupation_constraint = f"[{occupation}]"
+    else:
+        occupation_constraint = None
+
+    agent1_goal = st.text_input("Agent 1 Goal")
+    agent2_goal = st.text_input("Agent 2 Goal")
+
+    if st.button("Create Scenario"):
+        scenario_profile = EnvironmentProfileWrapper(
+            codename=codename,
+            source=source,
+            scenario=scenario,
+            relationship=relationship,
+            age_constraint=age_constraint,
+            occupation_constraint=occupation_constraint,
+            agent_goals=[agent1_goal, agent2_goal],
+        )
+
+        response = requests.post(
+            f"{st.session_state.API_BASE}/scenarios/",
+            json=scenario_profile.model_dump(),
+        )
+
+        if response.status_code != 200:
+            st.error("Failed to create scenario. Error: " + response.text)
+        else:
+            # there are quotes in the response
+            scenario_id = response.json()
+
+            st.success("Scenario created successfully! Scenario ID: " + scenario_id)
+            retrieved_scenario = requests.get(
+                f"{st.session_state.API_BASE}/scenarios/id/{scenario_id}"
+            )
+            st.write(retrieved_scenario.json())
+
+
+rendering_scenario_form()