feat: FastAPI Implementation of Sotopia Part Two (w websocket)

sotopia/server.py

@@ -1,7 +1,7 @@
 import asyncio
 import itertools
 import logging
-from typing import Literal, Sequence, Type
+from typing import Literal, Sequence, Type, AsyncGenerator, Union
 
 import gin
 import rich
@@ -25,7 +25,7 @@
     unweighted_aggregate_evaluate,
 )
 from sotopia.generation_utils.generate import LLM_Name, agenerate_script
-from sotopia.messages import AgentAction, Message, Observation
+from sotopia.messages import AgentAction, Message, Observation, SimpleMessage
 from sotopia.messages.message_classes import (
     ScriptBackground,
     ScriptEnvironmentResponse,
@@ -104,6 +104,12 @@ def run_sync_server(
     return messages
 
 
+def flatten_listed_messages(
+    messages: list[list[tuple[str, str, Message]]],
+) -> list[tuple[str, str, Message]]:
+    return list(itertools.chain.from_iterable(messages))
+
+
 @gin.configurable
 async def arun_one_episode(
     env: ParallelSotopiaEnv,
@@ -113,102 +119,125 @@ async def arun_one_episode(
     json_in_script: bool = False,
     tag: str | None = None,
     push_to_db: bool = False,
-) -> list[tuple[str, str, Message]]:
+    streaming: bool = False,
+) -> Union[
+    list[tuple[str, str, Message]],
+    AsyncGenerator[list[list[tuple[str, str, Message]]], None],
+]:
     agents = Agents({agent.agent_name: agent for agent in agent_list})
-    environment_messages = env.reset(agents=agents, omniscient=omniscient)
-    agents.reset()
-
-    messages: list[list[tuple[str, str, Message]]] = []
 
-    # Main Event Loop
-    done = False
-    messages.append(
-        [
-            ("Environment", agent_name, environment_messages[agent_name])
-            for agent_name in env.agents
-        ]
-    )
-    # set goal for agents
-    for index, agent_name in enumerate(env.agents):
-        agents[agent_name].goal = env.profile.agent_goals[index]
-    rewards: list[list[float]] = []
-    reasons: list[str] = []
-    while not done:
-        # gather agent messages
-        agent_messages: dict[str, AgentAction] = dict()
-        actions = await asyncio.gather(
-            *[
-                agents[agent_name].aact(environment_messages[agent_name])
-                for agent_name in env.agents
-            ]
-        )
-        if script_like:
-            # manually mask one message
-            agent_mask = env.action_mask
-            for idx in range(len(agent_mask)):
-                print("Current mask: ", agent_mask)
-                if agent_mask[idx] == 0:
-                    print("Action not taken: ", actions[idx])
-                    actions[idx] = AgentAction(action_type="none", argument="")
-                else:
-                    print("Current action taken: ", actions[idx])
-
-        # actions = cast(list[AgentAction], actions)
-        for idx, agent_name in enumerate(env.agents):
-            agent_messages[agent_name] = actions[idx]
-
-            messages[-1].append((agent_name, "Environment", agent_messages[agent_name]))
+    async def generate_messages() -> (
+        AsyncGenerator[list[list[tuple[str, str, Message]]], None]
+    ):
+        environment_messages = env.reset(agents=agents, omniscient=omniscient)
+        agents.reset()
+        messages: list[list[tuple[str, str, Message]]] = []
 
-        # send agent messages to environment
-        (
-            environment_messages,
-            rewards_in_turn,
-            terminated,
-            ___,
-            info,
-        ) = await env.astep(agent_messages)
+        # Main Event Loop
+        done = False
         messages.append(
             [
                 ("Environment", agent_name, environment_messages[agent_name])
                 for agent_name in env.agents
             ]
         )
-        # print("Environment message: ", environment_messages)
-        # exit(0)
-        rewards.append([rewards_in_turn[agent_name] for agent_name in env.agents])
-        reasons.append(
-            " ".join(info[agent_name]["comments"] for agent_name in env.agents)
+        yield messages
+
+        # set goal for agents
+        for index, agent_name in enumerate(env.agents):
+            agents[agent_name].goal = env.profile.agent_goals[index]
+        rewards: list[list[float]] = []
+        reasons: list[str] = []
+        while not done:
+            # gather agent messages
+            agent_messages: dict[str, AgentAction] = dict()
+            actions = await asyncio.gather(
+                *[
+                    agents[agent_name].aact(environment_messages[agent_name])
+                    for agent_name in env.agents
+                ]
+            )
+            if script_like:
+                # manually mask one message
+                agent_mask = env.action_mask
+                for idx in range(len(agent_mask)):
+                    if agent_mask[idx] == 0:
+                        actions[idx] = AgentAction(action_type="none", argument="")
+                    else:
+                        pass
+
+            # actions = cast(list[AgentAction], actions)
+            for idx, agent_name in enumerate(env.agents):
+                agent_messages[agent_name] = actions[idx]
+
+                messages[-1].append(
+                    (agent_name, "Environment", agent_messages[agent_name])
+                )
+
+            # send agent messages to environment
+            (
+                environment_messages,
+                rewards_in_turn,
+                terminated,
+                ___,
+                info,
+            ) = await env.astep(agent_messages)
+            messages.append(
+                [
+                    ("Environment", agent_name, environment_messages[agent_name])
+                    for agent_name in env.agents
+                ]
+            )
+
+            yield messages
+            rewards.append([rewards_in_turn[agent_name] for agent_name in env.agents])
+            reasons.append(
+                " ".join(info[agent_name]["comments"] for agent_name in env.agents)
+            )
+            done = all(terminated.values())
+
+        epilog = EpisodeLog(
+            environment=env.profile.pk,
+            agents=[agent.profile.pk for agent in agent_list],
+            tag=tag,
+            models=[env.model_name, agent_list[0].model_name, agent_list[1].model_name],
+            messages=[
+                [(m[0], m[1], m[2].to_natural_language()) for m in messages_in_turn]
+                for messages_in_turn in messages
+            ],
+            reasoning=info[env.agents[0]]["comments"],
+            rewards=[info[agent_name]["complete_rating"] for agent_name in env.agents],
+            rewards_prompt=info["rewards_prompt"]["overall_prompt"],
         )
-        done = all(terminated.values())
+        rich.print(epilog.rewards_prompt)
+        agent_profiles, conversation = epilog.render_for_humans()
+        for agent_profile in agent_profiles:
+            rich.print(agent_profile)
+        for message in conversation:
+            rich.print(message)
+
+        if streaming:
+            # yield the rewards and reasonings
+            messages.append(
+                [("Evaluation", "Rewards", SimpleMessage(message=str(epilog.rewards)))]
+            )
+            messages.append(
+                [("Evaluation", "Reasoning", SimpleMessage(message=epilog.reasoning))]
+            )
+            yield messages
 
-    # TODO: clean up this part
-    epilog = EpisodeLog(
-        environment=env.profile.pk,
-        agents=[agent.profile.pk for agent in agent_list],
-        tag=tag,
-        models=[env.model_name, agent_list[0].model_name, agent_list[1].model_name],
-        messages=[
-            [(m[0], m[1], m[2].to_natural_language()) for m in messages_in_turn]
-            for messages_in_turn in messages
-        ],
-        reasoning=info[env.agents[0]]["comments"],
-        rewards=[info[agent_name]["complete_rating"] for agent_name in env.agents],
-        rewards_prompt=info["rewards_prompt"]["overall_prompt"],
-    )
-    rich.print(epilog.rewards_prompt)
-    agent_profiles, conversation = epilog.render_for_humans()
-    for agent_profile in agent_profiles:
-        rich.print(agent_profile)
-    for message in conversation:
-        rich.print(message)
+        if push_to_db:
+            try:
+                epilog.save()
+            except Exception as e:
+                logging.error(f"Failed to save episode log: {e}")
 
-    if push_to_db:
-        try:
-            epilog.save()
-        except Exception as e:
-            logging.error(f"Failed to save episode log: {e}")
-    # flatten nested list messages
-    return list(itertools.chain(*messages))
+    if streaming:
+        return generate_messages()
+    else:
+        async for last_messages in generate_messages():
+            pass
+        return flatten_listed_messages(last_messages)
 
 
 @gin.configurable

sotopia/ui/README.md

@@ -47,6 +47,13 @@ returns:
 - agents: list[AgentProfile]
 
 
+#### GET /episodes
+
+Get all episodes.
+
+returns:
+- episodes: list[Episode]
+
 #### GET /episodes/?get_by={id|tag}/{episode_id|episode_tag}
 
 Get episode by episode_tag.
@@ -78,33 +85,31 @@ EnvironmentProfile
 returns:
 - scenario_id: str
 
-#### DELETE /agents/{agent_id}
+### Updating Data in the API Server
 
-Delete agent profile from the API server.
+#### PUT /agents/{agent_id}
+
+Update agent profile in the API server.
+Request Body:
+AgentProfile
 
 returns:
 - agent_id: str
 
-#### DELETE /scenarios/{scenario_id}
 
-Delete scenario profile from the API server.
+#### PUT /scenarios/{scenario_id}
+
+Update scenario profile in the API server.
+Request Body:
+EnvironmentProfile
 
 returns:
 - scenario_id: str
 
-
-### Error Code
-For RESTful APIs above we have the following error codes:
-| **Error Code** | **Description**                      |
-|-----------------|--------------------------------------|
-| **404**         | A resource is not found             |
-| **403**         | The query is not authorized         |
-| **500**         | Internal running error              |
-
 ### Initiating a new non-streaming simulation episode
 
 #### POST /episodes/
-[!] Currently not planning to implement
+
 ```python
 class SimulationEpisodeInitiation(BaseModel):
     scenario_id: str
@@ -147,14 +152,14 @@ returns:
 | Type | Direction   | Description |
 |-----------|--------|-------------|
 | SERVER_MSG | Server → Client | Standard message from server (payload: `messageForRendering` [here](https://github.com/sotopia-lab/sotopia-demo/blob/main/socialstream/rendering_utils.py) ) |
-| CLIENT_MSG | Client → Server | Standard message from client (payload: Currently not needed) |
-| ERROR      | Server → Client | Error notification (payload: `{"type": ERROR_TYPE, "description": DESC}`) |
+| CLIENT_MSG | Client → Server | Standard message from client (payload: TBD) |
+| ERROR      | Server → Client | Error notification (payload: TBD) |
 | START_SIM  | Client → Server | Initialize simulation (payload: `SimulationEpisodeInitialization`) |
 | END_SIM    | Client → Server | End simulation (payload: not needed) |
 | FINISH_SIM | Server → Client | Terminate simulation (payload: not needed) |
 
 
-**ERROR_TYPE**
+**Error Type**
 
 | Error Code | Description |
 |------------|-------------|
@@ -167,14 +172,27 @@ returns:
 | OTHER | Other unspecified errors |
 
 
-**Conversation Message From the Server**
-The server returns messages encapsulated in a structured format which is defined as follows:
 
-```python
-class MessageForRendering(TypedDict):
-    role: str # Specifies the origin of the message. Common values include "Background Info", "Environment", "{Agent Names}
-    type: str # Categorizes the nature of the message. Common types include: "comment", "said", "action"
-    content: str
+**Implementation plan**: Currently only support LLM-LLM simulation based on [this function](https://github.com/sotopia-lab/sotopia/blob/19d39e068c3bca9246fc366e5759414f62284f93/sotopia/server.py#L108).
+
+
+## An example to run simulation with the API
+
+**Get all scenarios**:
+```bash
+curl -X GET "http://localhost:8000/scenarios"
 ```
 
-**Implementation plan**: Currently only support LLM-LLM simulation based on [this function](https://github.com/sotopia-lab/sotopia/blob/19d39e068c3bca9246fc366e5759414f62284f93/sotopia/server.py#L108).
+Randomly select a scenario, e.g., `01HZRGTG1K4YQ2CBS9SNH28R9S`
+
+
+**Get all agents**:
+```bash
+curl -X GET "http://localhost:8000/agents"
+```
+
+Randomly select two agents, e.g., `01H5TNE5PE9RQGH86YM6MSWZMW` and `01H5TNE5PT06B3QPXJ65HHACV7`
+
+**Connecting to the websocket server**:
+
+@bugsz: Adding ur example here?