[autofix.ci] apply automated fixes

Author: autofix-ci[bot]

tests/api/test_turntaking_atomic.py

@@ -1,25 +1,20 @@
 # test_turntaking_atomic.py
 
 import pytest
-import asyncio
 from typing import List, Dict, Any, Optional, Union
 import logging
 from pydantic import BaseModel
 from sotopia.api.websocket_utils import (
     build_observation,
     get_env_agents,
-    WSMessageType,
 )
 from sotopia.agents import LLMAgent, Agents
 from sotopia.messages import Observation, AgentAction
-from sotopia.envs.parallel import ParallelSotopiaEnv
-from sotopia.envs.evaluators import (
-    RuleBasedTerminatedEvaluator,
-    EpisodeLLMEvaluator,
-    EvaluationForTwoAgents,
-    SotopiaDimensions,
+from sotopia.database import (
+    EnvironmentProfile,
+    AgentProfile,
+    EvaluationDimensionBuilder,
 )
-from sotopia.database import EnvironmentProfile, AgentProfile, EvaluationDimensionBuilder
 from sotopia.database.persistent_profile import RelationshipType
 
 
@@ -35,6 +30,7 @@
 # Dummy classes to simulate database objects
 # =============================================================================
 
+
 # Define DummyAgentProfile as a subclass of AgentProfile so that it is compatible.
 class DummyAgentProfile(AgentProfile):
     @classmethod
@@ -67,6 +63,7 @@ def construct_dummy(
         }
         return cls.model_validate(data)  # Using model_validate() for pydantic models
 
+
 # Define DummyEnvProfile with all required fields.
 class DummyEnvProfile:
     def __init__(self, pk: str) -> None:
@@ -75,32 +72,42 @@ def __init__(self, pk: str) -> None:
         # Force source to be a string (not None)
         self.source: str = ""
         # Use a non-None scenario.
-        self.scenario: str = "A concrete scenario description that meets the guidelines."
+        self.scenario: str = (
+            "A concrete scenario description that meets the guidelines."
+        )
         # For test consistency, use lowercase goals if tests expect them.
         self.agent_goals: List[str] = ["goal1", "goal2"]
-        self.relationship: RelationshipType = RelationshipType.stranger  
+        self.relationship: RelationshipType = RelationshipType.stranger
         # For fields that might be str but can be None in model, ensure a string is provided.
         self.age_constraint: str = ""
         self.occupation_constraint: str = ""
         self.agent_constraint: List[List[str]] = []
         self.tag: str = "test_tag"
 
+
 def fake_agent_get(agent_id: str) -> DummyAgentProfile:
     dummy: Dict[str, DummyAgentProfile] = {
-        "agent1": DummyAgentProfile.construct_dummy("agent1", first_name="John", last_name="Doe", age=30),
-        "agent2": DummyAgentProfile.construct_dummy("agent2", first_name="Jane", last_name="Doe", age=25),
+        "agent1": DummyAgentProfile.construct_dummy(
+            "agent1", first_name="John", last_name="Doe", age=30
+        ),
+        "agent2": DummyAgentProfile.construct_dummy(
+            "agent2", first_name="Jane", last_name="Doe", age=25
+        ),
     }
     if agent_id in dummy:
         return dummy[agent_id]
     raise Exception(f"AgentProfile with id {agent_id} not found")
 
+
 def fake_env_get(env_id: str) -> DummyEnvProfile:
     return DummyEnvProfile(env_id)
 
+
 # =============================================================================
 # Fixture for monkeypatching
 # =============================================================================
 
+
 @pytest.fixture
 def mp(monkeypatch: pytest.MonkeyPatch) -> pytest.MonkeyPatch:
     monkeypatch.setattr(EnvironmentProfile, "get", fake_env_get)
@@ -122,22 +129,27 @@ def patched_init(
     ) -> None:
         if agent_name is None and agent_profile is not None:
             agent_name = agent_profile.pk  # This sets the agent name to its pk.
-        original_init(self, agent_name, uuid_str, agent_profile, model_name, script_like)
+        original_init(
+            self, agent_name, uuid_str, agent_profile, model_name, script_like
+        )
+
     monkeypatch.setattr(LLMAgent, "__init__", patched_init)
     return monkeypatch
 
+
 # =============================================================================
 # Test for build_observation
 # =============================================================================
 
+
 def test_build_observation() -> None:
     """
     Test that build_observation returns an Observation with the correct
     last_turn, turn_number and available_actions.
     """
     conversation_history: List[Dict[str, str]] = [
         {"role": "client", "content": "Hello"},
-        {"role": "agent", "content": "Hi, how may I help you?"}
+        {"role": "agent", "content": "Hi, how may I help you?"},
     ]
     turn_number: int = len(conversation_history)
     obs: Observation = build_observation(turn_number, conversation_history)
@@ -151,16 +163,20 @@ def test_build_observation() -> None:
     ]
     assert obs.available_actions == expected_actions
 
+
 # =============================================================================
 # Test for get_env_agents
 # =============================================================================
 
+
 def test_get_env_agents(mp: pytest.MonkeyPatch) -> None:
     """
     Test that get_env_agents returns an environment, an agents dictionary keyed by
     agent names (which should be the dummy profile pks) and non-empty environment messages.
     """
-    env, agents, env_msgs = get_env_agents("env1", ["agent1", "agent2"], ["model1", "model2"], "eval_model", "dummy_list")
+    env, agents, env_msgs = get_env_agents(
+        "env1", ["agent1", "agent2"], ["model1", "model2"], "eval_model", "dummy_list"
+    )
     # Because our patched LLMAgent.__init__ uses agent_profile.pk as agent_name,
     # the keys should be "agent1" and "agent2".
     assert set(agents.keys()) == {"Jane Doe", "John Doe"}
@@ -169,60 +185,87 @@ def test_get_env_agents(mp: pytest.MonkeyPatch) -> None:
         assert agent.goal in ["goal1", "goal2"]
     assert isinstance(env_msgs, dict)
 
+
 # =============================================================================
 # Atomic test for process_turn
 # =============================================================================
 
+
 # Create a dummy agent by subclassing LLMAgent that returns a fixed AgentAction.
 class DummyAgent(LLMAgent):
     async def aact(self, obs: Observation) -> AgentAction:
         return AgentAction(action_type="speak", argument="dummy response")
 
+
 @pytest.fixture
 def dummy_simulator(mp: pytest.MonkeyPatch) -> Any:
     """
     Create a dummy simulator that mimics WebSocketSotopiaSimulator.
     It sets up a dummy Agents dictionary with one DummyAgent,
     a dummy environment with one goal, and conversation history seeded with an initial message.
     """
-    agents_instance: Agents = Agents({
-        "agent1": DummyAgent(agent_name="agent1", agent_profile=None, model_name="dummy")
-    })
+    agents_instance: Agents = Agents(
+        {
+            "agent1": DummyAgent(
+                agent_name="agent1", agent_profile=None, model_name="dummy"
+            )
+        }
+    )
+
     class DummyEnv:
         def __init__(self, goals: List[str]) -> None:
             self.agents: List[str] = list(agents_instance.keys())
-            self.profile = type("DummyProfile", (), {"agent_goals": goals, "pk": "env1"})
+            self.profile = type(
+                "DummyProfile", (), {"agent_goals": goals, "pk": "env1"}
+            )
+
     dummy_env = DummyEnv(["goal1"])
     dummy_msgs: Dict[str, Any] = {
-        "agent1": type("DummyObs", (), {"to_natural_language": lambda self: "initial message"})()
+        "agent1": type(
+            "DummyObs", (), {"to_natural_language": lambda self: "initial message"}
+        )()
     }
+
     class DummySimulator:
         def __init__(self) -> None:
             self.env: DummyEnv = dummy_env
             self.agents: Agents = agents_instance
             self.environment_messages: Dict[str, Any] = dummy_msgs
-            self.conversation_history: List[Dict[str, str]] = [{
-                "role": "environment",
-                "agent": "agent1",
-                "content": dummy_msgs["agent1"].to_natural_language()
-            }]
-        async def process_turn(self, client_data: Dict[str, str]) -> Dict[str, Union[int, str]]:
-            self.conversation_history.append({"role": "client", "content": client_data.get("content", "")})
+            self.conversation_history: List[Dict[str, str]] = [
+                {
+                    "role": "environment",
+                    "agent": "agent1",
+                    "content": dummy_msgs["agent1"].to_natural_language(),
+                }
+            ]
+
+        async def process_turn(
+            self, client_data: Dict[str, str]
+        ) -> Dict[str, Union[int, str]]:
+            self.conversation_history.append(
+                {"role": "client", "content": client_data.get("content", "")}
+            )
             agent_id: str | None = client_data.get("agent_id")
             if agent_id not in self.agents:
                 raise ValueError(f"Agent with id {agent_id} not found")
-            obs: Observation = build_observation(len(self.conversation_history), self.conversation_history)
+            obs: Observation = build_observation(
+                len(self.conversation_history), self.conversation_history
+            )
             agent = self.agents[agent_id]
             agent_action: AgentAction = await agent.aact(obs)
-            self.conversation_history.append({"role": "agent", "content": agent_action.argument})
+            self.conversation_history.append(
+                {"role": "agent", "content": agent_action.argument}
+            )
             return {
                 "turn": len(self.conversation_history),
                 "agent_id": agent_id,
                 "agent_response": agent_action.argument,
                 "action_type": agent_action.action_type,
             }
+
     return DummySimulator()
 
+
 @pytest.mark.asyncio
 async def test_process_turn_success(dummy_simulator: Any) -> None:
     """
@@ -237,6 +280,7 @@ async def test_process_turn_success(dummy_simulator: Any) -> None:
     assert result["agent_response"] == "dummy response"
     assert result["action_type"] == "speak"
 
+
 @pytest.mark.asyncio
 async def test_process_turn_invalid_agent(dummy_simulator: Any) -> None:
     """
@@ -248,6 +292,7 @@ async def test_process_turn_invalid_agent(dummy_simulator: Any) -> None:
         await simulator.process_turn(client_data)
     assert "Agent with id nonexistent not found" in str(excinfo.value)
 
+
 @pytest.mark.asyncio
 async def test_multiple_turns_accumulate_history(dummy_simulator: Any) -> None:
     """