sotopia-lab · ProKil · Nov 22, 2024 · Oct 30, 2024 · Nov 5, 2024 · Nov 7, 2024
diff --git a/pyproject.toml b/pyproject.toml
@@ -37,7 +37,10 @@ groq = ["groq"]
 cohere = ["cohere"]
 google-generativeai = ["google-generativeai"]
 examples = ["transformers", "datasets", "scipy", "torch", "pandas"]
-chat = ["fastapi"]
+chat = [
+    "fastapi",
+    "websockets>=13.1",
+]
 test = ["pytest", "pytest-cov", "pytest-asyncio"]
 
 [tool.uv]

diff --git a/sotopia/ui/README.md b/sotopia/ui/README.md
@@ -0,0 +1,162 @@
+# Sotopia UI
+
+## FastAPI Server
+
+The API server is a FastAPI application that is used to connect the Sotopia UI to the Sotopia backend.
+This could also help with other projects that need to connect to the Sotopia backend through HTTP requests.
+
+Here are some initial design of the API server:
+
+### Getting Data from the API Server
+
+#### GET /get/scenarios/{scenario_id}
+
+Get scenarios by scenario_id.
+parameters:
+- scenario_id: str
+
+returns:
+- scenarios: EnvironmentProfile
+
+
+#### GET /get/scenarios
+
+Get all scenarios.
+
+returns:
+- scenarios: list[EnvironmentProfile]
+
+#### GET /get/scenarios/{sceanrio_tag}
+
+Get scenarios by scenario_tag.
+parameters:
+- scenario_tag: str
+(This scenario tag could be a keyword; so people can search for scenarios by keywords)
+
+returns:
+- scenarios: list[EnvironmentProfile]
+
+#### GET /get/agents
+
+Get all agents.
+
+returns:
+- agents: list[AgentProfile]
+
+#### GET /get/agents/{agent_id}
+
+Get agent by agent_id.
+parameters:
+- agent_id: str
+
+returns:
+- agent: AgentProfile
+
+#### GET /get/agents/{agent_gender}
+
+Get agents by agent_gender.
+parameters:
+- agent_gender: Literal["male", "female"]
+
+returns:
+- agents: list[AgentProfile]
+
+#### GET /get/agents/{agent_occupation}
+
+Get agents by agent_occupation.
+parameters:
+- agent_occupation: str
+
+returns:
+- agents: list[AgentProfile]
+
+
+#### GET /get/episodes
+
+Get all episodes.
+
+returns:
+- episodes: list[Episode]
+
+#### GET /get/episodes/{episode_tag}
+
+Get episode by episode_tag.
+parameters:
+- episode_tag: str
+
+returns:
+- episode: list[Episode]
+
+#### GET /get/episodes/{episode_id}
+
+Get episode by episode_id.
+parameters:
+- episode_id: str
+
+returns:
+- episode: Episode
+
+
+### Sending Data to the API Server
+
+w#### POST /post/agents/
+
+Send agent profile to the API server.
+Request Body:
+AgentProfile
+
+returns:
+- agent_id: str
+
+#### POST /post/scenarios/
+
+Send scenario profile to the API server.
+Request Body:
+EnvironmentProfile
+
+returns:
+- scenario_id: str
+
+### Updating Data in the API Server
+
+#### PUT /put/agents/{agent_id}
+
+Update agent profile in the API server.
+Request Body:
+AgentProfile
+
+returns:
+- agent_id: str
+
+
+#### PUT /put/scenarios/{scenario_id}
+
+Update scenario profile in the API server.
+Request Body:
+EnvironmentProfile
+
+returns:
+- scenario_id: str
+
+### Initiating a new non-streaming simulation episode
+
+#### POST /post/episodes/
+
+```python
+class SimulationEpisodeInitiation(BaseModel):
+    scenario_id: str
+    agent_ids: list[str]
+    episode_tag: str
+    models: list[str]
+```
+
+Send episode profile to the API server.
+Request Body:
+SimulationEpisodeInitiation
+
+returns:
+- episode_id: str (This is the id of the episode that will be used to get the episode data, saved in the redis database)
+
+### Initiating a new interactive streaming simulation episode (this operation will open a websocket connection)
+
+We use the websocket connection to send the simulation step-by-step results to the UI.
diff --git a/sotopia/ui/websocket_ex.py b/sotopia/ui/websocket_ex.py
@@ -0,0 +1,81 @@
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.responses import HTMLResponse
+
+app = FastAPI()
+
+html = """
+<!DOCTYPE html>
+<html>
+    <head>
+        <title>Chat</title>
+    </head>
+    <body>
+        <h1>WebSocket Chat</h1>
+        <h2>Your ID: <span id="ws-id"></span></h2>
+        <form action="" onsubmit="sendMessage(event)">
+            <input type="text" id="messageText" autocomplete="off"/>
+            <button>Send</button>
+        </form>
+        <ul id='messages'>
+        </ul>
+        <script>
+            var client_id = Date.now()
+            document.querySelector("#ws-id").textContent = client_id;
+            var ws = new WebSocket(`ws://localhost:8000/ws/${client_id}`);
+            ws.onmessage = function(event) {
+                var messages = document.getElementById('messages')
+                var message = document.createElement('li')
+                var content = document.createTextNode(event.data)
+                message.appendChild(content)
+                messages.appendChild(message)
+            };
+            function sendMessage(event) {
+                var input = document.getElementById("messageText")
+                ws.send(input.value)
+                input.value = ''
+                event.preventDefault()
+            }
+        </script>
+    </body>
+</html>
+"""
+
+
+class ConnectionManager:
+    def __init__(self) -> None:
+        self.active_connections: list[WebSocket] = []
+
+    async def connect(self, websocket: WebSocket) -> None:
+        await websocket.accept()
+        self.active_connections.append(websocket)
+
+    def disconnect(self, websocket: WebSocket) -> None:
+        self.active_connections.remove(websocket)
+
+    async def send_personal_message(self, message: str, websocket: WebSocket) -> None:
+        await websocket.send_text(message)
+
+    async def broadcast(self, message: str) -> None:
+        for connection in self.active_connections:
+            await connection.send_text(message)
+
+
+manager = ConnectionManager()
+
+
+@app.get("/")
+async def get() -> HTMLResponse:
+    return HTMLResponse(html)
+
+
+@app.websocket("/ws/{client_id}")
+async def websocket_endpoint(websocket: WebSocket, client_id: int) -> None:
+    await manager.connect(websocket)
+    try:
+        while True:
+            data = await websocket.receive_text()
+            await manager.send_personal_message(f"You wrote: {data}", websocket)
+            await manager.broadcast(f"Client #{client_id} says: {data}")
+    except WebSocketDisconnect:
+        manager.disconnect(websocket)
+        await manager.broadcast(f"Client #{client_id} left the chat")