From cfbb2334603d07bbece45e48c56704f4695adc8c Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Wed, 16 Apr 2025 12:43:01 -0400
Subject: [PATCH 01/23] custom agent

---
 sotopia/cli/benchmark/benchmark.py | 23 +++++++++++++++++------
 sotopia/database/logs.py           |  1 +
 sotopia/server.py                  |  3 ++-
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index 401fee3bd..172f13e9e 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -329,6 +329,7 @@ def _list_all_env_agent_combo_not_in_db(
     env_agent_combo_storage_index_list: list[tuple[EnvAgentComboStorage, str]],
     tag: str = "",
     task: str = "",
+    agent_class: type[LLMAgent] = LLMAgent,
 ) -> list[EnvAgentCombo[Observation, AgentAction]]:
     """Iterate over each environment and return the first env-agent combo not in the database."""
     assert tag, "tag should not be empty"
@@ -370,12 +371,16 @@ def _list_all_env_agent_combo_not_in_db(
         agent_profiles = [AgentProfile.get(id) for id in agent_ids]
         # make sure the second agent (i.e., the agent being benchmarked) is always the indexed agent
         agents = [
-            LLMAgent(agent_profile=agent_profile, model_name=agent_model)
-            for agent_profile, agent_model in zip(
-                agent_profiles,
-                [model_names["test_model"], model_names["partner_model"]]
-                if index == "0"
-                else [model_names["partner_model"], model_names["test_model"]],
+            agent_class(agent_profile=agent_profile, model_name=agent_model)
+            if (index == "0" and i == 0) or (index == "1" and i == 1)
+            else LLMAgent(agent_profile=agent_profile, model_name=agent_model)
+            for i, (agent_profile, agent_model) in enumerate(
+                zip(
+                    agent_profiles,
+                    [model_names["test_model"], model_names["partner_model"]]
+                    if index == "0"
+                    else [model_names["partner_model"], model_names["test_model"]],
+                )
             )
         ]
         list_of_env_agent_combo_storage.append((env, agents))
@@ -435,6 +440,10 @@ def benchmark(
         default_model_list,
         help=f"All the language model you want to benchmark. Default is the pre-loaded model list {default_model_list}.",
     ),
+    agent_class: type[LLMAgent] = typer.Option(
+        LLMAgent,
+        help="The agent class you want to use for the evaluated models. Must be a subclass of LLMAgent.",
+    ),
     partner_model: str = typer.Option(
         "together_ai/meta-llama/Llama-3-70b-chat-hf",
         help="The partner model you want to use.",
@@ -496,6 +505,7 @@ def benchmark(
             tag=tag,
             env_agent_combo_storage_index_list=env_agent_combo_storage_index_list,
             task=task,
+            agent_class=agent_class,
         )
         number_of_fix_turns = 0
         while True:
@@ -511,6 +521,7 @@ def benchmark(
                 tag=tag,
                 env_agent_combo_storage_index_list=env_agent_combo_storage_index_list,
                 task=task,
+                agent_class=agent_class,
             )
             number_of_fix_turns += 1
             if len(env_agent_combo_list) == 0 or number_of_fix_turns >= 5:
diff --git a/sotopia/database/logs.py b/sotopia/database/logs.py
index da66c7fcf..f30b34f7d 100644
--- a/sotopia/database/logs.py
+++ b/sotopia/database/logs.py
@@ -26,6 +26,7 @@ class BaseEpisodeLog(BaseModel):
     agents: list[str] = Field(index=True)
     tag: str | None = Field(index=True, default="")
     models: list[str] | None = Field(index=True, default=[])
+    agent_classes: list[str] | None = Field(index=True, default=[])
     messages: list[list[tuple[str, str, str]]]  # Messages arranged by turn
     reasoning: str = Field(default="")
     rewards: list[tuple[float, dict[str, float]] | float]  # Rewards arranged by turn
diff --git a/sotopia/server.py b/sotopia/server.py
index 9e8e76875..72bdc19e5 100644
--- a/sotopia/server.py
+++ b/sotopia/server.py
@@ -195,12 +195,13 @@ async def generate_messages() -> (
                 " ".join(info[agent_name]["comments"] for agent_name in env.agents)
             )
             done = all(terminated.values())
-
+        breakpoint()
         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],
+            agent_classes=[type(agent).__name__ for agent in agent_list],
             messages=[
                 [(m[0], m[1], m[2].to_natural_language()) for m in messages_in_turn]
                 for messages_in_turn in messages

From 3efe458580c55b8826213b65cf58286eceaba2a8 Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Wed, 16 Apr 2025 13:01:36 -0400
Subject: [PATCH 02/23] self-defined tag option

---
 sotopia/cli/benchmark/benchmark.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index 172f13e9e..74995cd1e 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -459,6 +459,7 @@ def benchmark(
     output_to_jsonl: bool = typer.Option(False, help="Output to jsonl."),
     push_to_db: bool = typer.Option(False, help="Push to db."),
     save_dir: str = typer.Option(".", help="The directory to save the output."),
+    tag: str = typer.Option("", help="The tag for the experiment."),
 ) -> None:
     if only_show_performance:
         benchmark_display(
@@ -493,7 +494,11 @@ def benchmark(
         typer.echo(
             f"Running benchmark for {model} chatting with {partner_model} on task {task} with {evaluator_model} as the evaluator."
         )
-        tag = f"benchmark_{model}_{partner_model}_{evaluator_model}_{task}_trial0"
+        tag = (
+            f"benchmark_{model}_{partner_model}_{evaluator_model}_{task}_trial0"
+            if tag == ""
+            else tag
+        )
         typer.echo(typer.style(f"Tag: {tag}", fg=typer.colors.GREEN, bold=True))
         model_names = {
             "env": evaluator_model,

From 2343097e411fb93e0a8e4ae05a5f675dda0b7587 Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Wed, 16 Apr 2025 13:53:32 -0400
Subject: [PATCH 03/23] better cli

---
 sotopia/cli/benchmark/benchmark.py | 62 ++++++++++++++++++------------
 1 file changed, 38 insertions(+), 24 deletions(-)

diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index 74995cd1e..1557932ee 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -42,6 +42,7 @@
 from pathlib import Path
 from ..app import app
 import os
+from typing import Annotated
 
 # date and message only
 FORMAT = "%(asctime)s - %(levelname)s - %(name)s - %(message)s"
@@ -436,30 +437,43 @@ def save_to_jsonl(
 
 @app.command()
 def benchmark(
-    models: list[str] = typer.Option(
-        default_model_list,
-        help=f"All the language model you want to benchmark. Default is the pre-loaded model list {default_model_list}.",
-    ),
-    agent_class: type[LLMAgent] = typer.Option(
-        LLMAgent,
-        help="The agent class you want to use for the evaluated models. Must be a subclass of LLMAgent.",
-    ),
-    partner_model: str = typer.Option(
-        "together_ai/meta-llama/Llama-3-70b-chat-hf",
-        help="The partner model you want to use.",
-    ),
-    evaluator_model: str = typer.Option(
-        "gpt-4o", help="The evaluator model you want to use."
-    ),
-    batch_size: int = typer.Option(10, help="The batch size you want to use."),
-    task: str = typer.Option("hard", help="The task id you want to benchmark."),
-    url: str = typer.Option("", help="The url to fetch the benchmark combo."),
-    print_logs: bool = typer.Option(False, help="Print logs."),
-    only_show_performance: bool = typer.Option(False, help="Only show performance."),
-    output_to_jsonl: bool = typer.Option(False, help="Output to jsonl."),
-    push_to_db: bool = typer.Option(False, help="Push to db."),
-    save_dir: str = typer.Option(".", help="The directory to save the output."),
-    tag: str = typer.Option("", help="The tag for the experiment."),
+    models: Annotated[
+        list[str],
+        typer.Argument(
+            help=f"All the language model you want to benchmark. Default is the pre-loaded model list {default_model_list}."
+        ),
+    ] = default_model_list,
+    agent_class: Annotated[
+        type[LLMAgent],
+        typer.Argument(
+            help="The agent class you want to use for the evaluated models. Must be a subclass of LLMAgent."
+        ),
+    ] = LLMAgent,
+    partner_model: Annotated[
+        str, typer.Argument(help="The partner model you want to use.")
+    ] = "together_ai/meta-llama/Llama-3-70b-chat-hf",
+    evaluator_model: Annotated[
+        str, typer.Argument(help="The evaluator model you want to use.")
+    ] = "gpt-4o",
+    batch_size: Annotated[
+        int, typer.Argument(help="The batch size you want to use.")
+    ] = 10,
+    task: Annotated[
+        str, typer.Argument(help="The task id you want to benchmark.")
+    ] = "hard",
+    url: Annotated[
+        str, typer.Argument(help="The url to fetch the benchmark combo.")
+    ] = "",
+    print_logs: Annotated[bool, typer.Argument(help="Print logs.")] = False,
+    only_show_performance: Annotated[
+        bool, typer.Argument(help="Only show performance.")
+    ] = False,
+    output_to_jsonl: Annotated[bool, typer.Argument(help="Output to jsonl.")] = False,
+    push_to_db: Annotated[bool, typer.Argument(help="Push to db.")] = False,
+    save_dir: Annotated[
+        str, typer.Argument(help="The directory to save the output.")
+    ] = ".",
+    tag: Annotated[str, typer.Argument(help="The tag for the experiment.")] = "",
 ) -> None:
     if only_show_performance:
         benchmark_display(

From 0cd503a6e79f20422b088448deeafc0472407b2c Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Wed, 16 Apr 2025 14:40:05 -0400
Subject: [PATCH 04/23] update checking mechenism

---
 sotopia/cli/benchmark/benchmark.py | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index 1557932ee..71dc851f4 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -206,16 +206,18 @@ def _set_up_logs(
 
 def preprocess_episode_data(
     episode_list_with_tag: list[EpisodeLog],
-) -> dict[tuple[str, tuple[str, ...], tuple[str, ...]], bool]:
+) -> dict[tuple[str, tuple[str, ...], tuple[str, ...], tuple[str, str]], bool]:
     """Preprocess episodes into a dictionary for quick lookup."""
     episode_dict = {}
     for episode in episode_list_with_tag:
         assert isinstance(episode, EpisodeLog), "episode should be an EpisodeLog"
         assert episode.models is not None, "episode.models should not be None"
+        agent_classes = getattr(episode, "agent_classes", None)
         episode_key = (
             episode.environment,
             tuple(episode.agents),
             tuple(episode.models),
+            tuple(agent_classes) if agent_classes else ("", ""),
         )
         episode_dict[episode_key] = True
     return episode_dict
@@ -226,14 +228,25 @@ def check_existing_episodes(
     agent_ids: list[str],
     models: dict[str, str],
     index: str,
-    episode_dict: dict[tuple[str, tuple[str, ...], tuple[str, ...]], bool],
+    episode_dict: dict[
+        tuple[str, tuple[str, ...], tuple[str, ...], tuple[str, str]], bool
+    ],
+    agent_classes: dict[str, str],
 ) -> bool:
     models_list = (
         (models["env"], models["test_model"], models["partner_model"])
         if index == "0"
         else (models["env"], models["partner_model"], models["test_model"])
     )
-    episode_key = (env_id, tuple(agent_ids), models_list)
+    if agent_classes["test_model"] != agent_classes["partner_model"]:
+        agent_classes_list = (
+            (agent_classes["test_model"], agent_classes["partner_model"])
+            if index == "0"
+            else (agent_classes["partner_model"], agent_classes["test_model"])
+        )
+    else:
+        agent_classes_list = ("", "")
+    episode_key = (env_id, tuple(agent_ids), models_list, agent_classes_list)
     return episode_dict.get(episode_key, False)
 
 
@@ -350,6 +363,10 @@ def _list_all_env_agent_combo_not_in_db(
             models=model_names,
             index=index,
             episode_dict=episode_dict,
+            agent_classes={
+                "test_model": agent_class.__name__,
+                "partner_model": LLMAgent.__name__,
+            },
         ):
             logging.info(
                 f"Episode for {env_id} with agents {agent_ids} using {list(model_names.values())} already exists"

From b0bf1f10abff66ab2f1cb8c78bab1ed8f5d338cb Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Wed, 16 Apr 2025 19:55:55 -0400
Subject: [PATCH 05/23] rewards display

---
 sotopia/cli/benchmark/benchmark.py | 123 +++++++++++++++++++++++------
 1 file changed, 99 insertions(+), 24 deletions(-)

diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index 71dc851f4..140fe8954 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -73,10 +73,13 @@
 
 
 def get_avg_reward(
-    episodes: list[EpisodeLog], model_name: str
+    episodes: list[EpisodeLog],
+    model_name: str,
+    agent_class: str = "",
+    for_partner: bool = False,
 ) -> dict[str, tuple[float, float]]:
     """
-    input: list of EpisodeLog, model_name
+    input: list of EpisodeLog, model_name, agent_class
 
     return: dictionary of {dimension: (avg_reward, margin_of_error (in 95% confidence interval))}, plus the distinct setting number and episode count (in the same format, but with 0 margin of error)
     """
@@ -86,9 +89,16 @@ def get_avg_reward(
     avg_reward_dict = {}
     avg_margin_dict = {}
     avg_variance_dict = {}
+    if for_partner:
+        index = 1
+    else:
+        index = 0
     for episode in episodes:
         assert episode.models is not None, "episode.models should not be None"
-        if episode.models[1] == model_name:
+        agent_classes = getattr(episode, "agent_classes", None)
+        if episode.models[index + 1] == model_name and (
+            not agent_class or (agent_classes and agent_classes[index] == agent_class)
+        ):
             reward = get_rewards_from_episode(episode)[0][1]
             rewards_dict[f"{episode.environment}_0"].append(reward)
         else:
@@ -311,6 +321,7 @@ def benchmark_display(
     task: str = "hard",
     output_to_jsonl: bool = False,
     save_dir: str = ".",
+    agent_class: str = "",
 ) -> dict[str, dict[str, tuple[float, float]]]:
     """
     Usage: sotopia benchmark-display --model-list gpt-4o --model-list together_ai/meta-llama-Llama-3-70b-chat-hf
@@ -325,9 +336,14 @@ def benchmark_display(
         if len(episodes) == 0:
             print(f"No episodes found for {model}")
             continue
-        avg_rewards = get_avg_reward(episodes, model)  # type: ignore
-        model_rewards_dict[model] = avg_rewards
-        # print(f"Model: {model}, episodes: {len(episodes)}, Avg Rewards: {avg_rewards}")
+        # Get test model results
+        test_model_rewards = get_avg_reward(episodes, model, agent_class)  # type: ignore
+        # Get partner model results
+        partner_model_rewards = get_avg_reward(
+            episodes, partner_model, agent_class, for_partner=True
+        )  # type: ignore
+        model_rewards_dict[f"{model} (test)"] = test_model_rewards
+        model_rewards_dict[f"{partner_model} (partner)"] = partner_model_rewards
         rich.print(model_rewards_dict)
     if model_rewards_dict:
         display_in_table(model_rewards_dict)
@@ -408,23 +424,71 @@ def _list_all_env_agent_combo_not_in_db(
 def display_in_table(
     model_rewards_dict: dict[str, dict[str, tuple[float, float]]],
 ) -> None:
-    table = rich.table.Table(
-        title="Rewards Evaluation (+/- CI bounds)",
-    )
-    table.add_column("Model")
-    for dimension in model_rewards_dict[list(model_rewards_dict.keys())[0]].keys():
-        table.add_column(dimension)
-    table.add_column("Settings")
-    table.add_column("Episodes")
-    for model, rewards in model_rewards_dict.items():
-        table.add_row(
-            model,
-            *(
-                [f"{rewards[k][0]:.2f} ± {rewards[k][1]:.2f}" for k in rewards.keys()]
-                + [f"{rewards[k][0]:.0f}" for k in ["setting_num", "episode_count"]]
-            ),
+    # Group models by their type (test/partner)
+    test_models = {k: v for k, v in model_rewards_dict.items() if "(test)" in k}
+    partner_models = {k: v for k, v in model_rewards_dict.items() if "(partner)" in k}
+
+    # Create a table for test models
+    if test_models:
+        test_table = rich.table.Table(
+            title="Test Model Performance (+/- CI bounds)",
+            show_header=True,
+            header_style="bold magenta",
         )
-    rich.print(table)
+        test_table.add_column("Model")
+        for dimension in list(test_models.values())[0].keys():
+            if dimension not in ["setting_num", "episode_count"]:
+                test_table.add_column(dimension)
+        test_table.add_column("Settings")
+        test_table.add_column("Episodes")
+
+        for model, rewards in test_models.items():
+            test_table.add_row(
+                model.replace(" (test)", ""),
+                *(
+                    [
+                        f"{rewards[k][0]:.2f} ± {rewards[k][1]:.2f}"
+                        for k in rewards.keys()
+                        if k not in ["setting_num", "episode_count"]
+                    ]
+                    + [
+                        f"{rewards['setting_num'][0]:.0f}",
+                        f"{rewards['episode_count'][0]:.0f}",
+                    ]
+                ),
+            )
+        rich.print(test_table)
+
+    # Create a table for partner models
+    if partner_models:
+        partner_table = rich.table.Table(
+            title="Partner Model Performance (+/- CI bounds)",
+            show_header=True,
+            header_style="bold cyan",
+        )
+        partner_table.add_column("Model")
+        for dimension in list(partner_models.values())[0].keys():
+            if dimension not in ["setting_num", "episode_count"]:
+                partner_table.add_column(dimension)
+        partner_table.add_column("Settings")
+        partner_table.add_column("Episodes")
+
+        for model, rewards in partner_models.items():
+            partner_table.add_row(
+                model.replace(" (partner)", ""),
+                *(
+                    [
+                        f"{rewards[k][0]:.2f} ± {rewards[k][1]:.2f}"
+                        for k in rewards.keys()
+                        if k not in ["setting_num", "episode_count"]
+                    ]
+                    + [
+                        f"{rewards['setting_num'][0]:.0f}",
+                        f"{rewards['episode_count'][0]:.0f}",
+                    ]
+                ),
+            )
+        rich.print(partner_table)
 
 
 def save_to_jsonl(
@@ -494,7 +558,13 @@ def benchmark(
 ) -> None:
     if only_show_performance:
         benchmark_display(
-            models, partner_model, evaluator_model, task, output_to_jsonl, save_dir
+            models,
+            partner_model,
+            evaluator_model,
+            task,
+            output_to_jsonl,
+            save_dir,
+            agent_class.__name__,
         )
         return
 
@@ -564,5 +634,10 @@ def benchmark(
                 return
 
     benchmark_display(
-        models, partner_model, evaluator_model, task, output_to_jsonl=output_to_jsonl
+        models,
+        partner_model,
+        evaluator_model,
+        task,
+        output_to_jsonl=output_to_jsonl,
+        agent_class=agent_class.__name__,
     )

From 85291dcd6c2cc530a5bff3da45cb582ed354f24e Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Wed, 16 Apr 2025 20:06:52 -0400
Subject: [PATCH 06/23] agenerate action structure output

---
 sotopia/cli/benchmark/benchmark.py   | 7 +++++--
 sotopia/generation_utils/generate.py | 2 ++
 sotopia/server.py                    | 1 -
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index 140fe8954..242b40cb9 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -340,8 +340,11 @@ def benchmark_display(
         test_model_rewards = get_avg_reward(episodes, model, agent_class)  # type: ignore
         # Get partner model results
         partner_model_rewards = get_avg_reward(
-            episodes, partner_model, agent_class, for_partner=True
-        )  # type: ignore
+            episodes,
+            partner_model,
+            agent_class,
+            for_partner=True,  # type: ignore
+        )
         model_rewards_dict[f"{model} (test)"] = test_model_rewards
         model_rewards_dict[f"{partner_model} (partner)"] = partner_model_rewards
         rich.print(model_rewards_dict)
diff --git a/sotopia/generation_utils/generate.py b/sotopia/generation_utils/generate.py
index 21b4358b8..ce5e6fe76 100644
--- a/sotopia/generation_utils/generate.py
+++ b/sotopia/generation_utils/generate.py
@@ -254,6 +254,7 @@ async def agenerate_action(
     script_like: bool = False,
     bad_output_process_model: str | None = None,
     use_fixed_model_version: bool = True,
+    structured_output: bool = False,
 ) -> AgentAction:
     """
     Using langchain to generate an example episode
@@ -305,6 +306,7 @@ async def agenerate_action(
             temperature=temperature,
             bad_output_process_model=bad_output_process_model,
             use_fixed_model_version=use_fixed_model_version,
+            structured_output=structured_output,
         )
     except Exception as e:
         log.warning(f"Failed to generate action due to {e}")
diff --git a/sotopia/server.py b/sotopia/server.py
index 72bdc19e5..61a6d5f4d 100644
--- a/sotopia/server.py
+++ b/sotopia/server.py
@@ -195,7 +195,6 @@ async def generate_messages() -> (
                 " ".join(info[agent_name]["comments"] for agent_name in env.agents)
             )
             done = all(terminated.values())
-        breakpoint()
         epilog = EpisodeLog(
             environment=env.profile.pk,
             agents=[agent.profile.pk for agent in agent_list],

From 8b4d83da0f2e00d3a876100414f431944944f1e8 Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Wed, 16 Apr 2025 20:25:08 -0400
Subject: [PATCH 07/23] update environmentlist model

---
 sotopia/database/persistent_profile.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/sotopia/database/persistent_profile.py b/sotopia/database/persistent_profile.py
index f570c2221..5917208eb 100644
--- a/sotopia/database/persistent_profile.py
+++ b/sotopia/database/persistent_profile.py
@@ -116,7 +116,6 @@ class RelationshipProfile(BaseRelationshipProfile, JsonModel):
 
 
 class EnvironmentList(JsonModel):
-    pk: str = Field(default_factory=lambda: "")
     name: str = Field(index=True)
     environments: list[str] = Field(default_factory=lambda: [])
     agent_index: list[str] | None = Field(default_factory=lambda: None)

From 7fc26121e0ad68f1f5007a9f2d63f0653e9ffbbb Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Thu, 17 Apr 2025 00:49:25 -0400
Subject: [PATCH 08/23] update env name

---
 sotopia/cli/benchmark/benchmark.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index 242b40cb9..111fe631e 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -340,10 +340,10 @@ def benchmark_display(
         test_model_rewards = get_avg_reward(episodes, model, agent_class)  # type: ignore
         # Get partner model results
         partner_model_rewards = get_avg_reward(
-            episodes,
+            episodes,  # type: ignore
             partner_model,
             agent_class,
-            for_partner=True,  # type: ignore
+            for_partner=True,
         )
         model_rewards_dict[f"{model} (test)"] = test_model_rewards
         model_rewards_dict[f"{partner_model} (partner)"] = partner_model_rewards
@@ -395,6 +395,7 @@ def _list_all_env_agent_combo_not_in_db(
         env = ParallelSotopiaEnv(
             env_profile=env_profile,
             action_order="round-robin",
+            model_name=model_names["env"],
             evaluators=[
                 RuleBasedTerminatedEvaluator(max_turn_number=20, max_stale_turn=2),
             ],

From 620e3d3fa832379d751b0c0afeac0a22e2249827 Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Thu, 17 Apr 2025 12:53:01 -0400
Subject: [PATCH 09/23] benchmark display update

---
 sotopia/cli/benchmark/benchmark.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index 111fe631e..c58997a13 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -635,7 +635,7 @@ def benchmark(
             )
             number_of_fix_turns += 1
             if len(env_agent_combo_list) == 0 or number_of_fix_turns >= 5:
-                return
+                break
 
     benchmark_display(
         models,

From 23e33325158cf33c9bbf2202b0848339b4518c74 Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Thu, 17 Apr 2025 13:03:53 -0400
Subject: [PATCH 10/23] benchmark refine

---
 sotopia/cli/benchmark/benchmark.py | 31 ++++++++++--------------------
 1 file changed, 10 insertions(+), 21 deletions(-)

diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index c58997a13..dc8c39523 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -322,6 +322,7 @@ def benchmark_display(
     output_to_jsonl: bool = False,
     save_dir: str = ".",
     agent_class: str = "",
+    tag: str = "",
 ) -> dict[str, dict[str, tuple[float, float]]]:
     """
     Usage: sotopia benchmark-display --model-list gpt-4o --model-list together_ai/meta-llama-Llama-3-70b-chat-hf
@@ -331,7 +332,6 @@ def benchmark_display(
     print(f"Displaying performance for {model_list} vs {partner_model} on task {task}")
     model_rewards_dict = dict()
     for model in model_list:
-        tag = f"benchmark_{model}_{partner_model}_{evaluator_model}_{task}_trial0"
         episodes = EpisodeLog.find(EpisodeLog.tag == tag).all()
         if len(episodes) == 0:
             print(f"No episodes found for {model}")
@@ -560,18 +560,6 @@ def benchmark(
     ] = ".",
     tag: Annotated[str, typer.Argument(help="The tag for the experiment.")] = "",
 ) -> None:
-    if only_show_performance:
-        benchmark_display(
-            models,
-            partner_model,
-            evaluator_model,
-            task,
-            output_to_jsonl,
-            save_dir,
-            agent_class.__name__,
-        )
-        return
-
     """A simple command-line interface example."""
     _set_up_logs(print_logs=print_logs)
     benchmark_combo = initialize_benchmark_combo(url)
@@ -637,11 +625,12 @@ def benchmark(
             if len(env_agent_combo_list) == 0 or number_of_fix_turns >= 5:
                 break
 
-    benchmark_display(
-        models,
-        partner_model,
-        evaluator_model,
-        task,
-        output_to_jsonl=output_to_jsonl,
-        agent_class=agent_class.__name__,
-    )
+        benchmark_display(
+            [model],
+            partner_model,
+            evaluator_model,
+            task,
+            output_to_jsonl=output_to_jsonl,
+            agent_class=agent_class.__name__,
+            tag=tag,
+        )

From 7766c2e19d46af0dd0594da09b87042e80771f47 Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Thu, 17 Apr 2025 17:45:15 -0400
Subject: [PATCH 11/23] force llm agent structured output

---
 sotopia/agents/llm_agent.py        |  1 +
 sotopia/cli/benchmark/benchmark.py | 28 +++++++++++++++++++---------
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/sotopia/agents/llm_agent.py b/sotopia/agents/llm_agent.py
index 497954d7b..3269562cd 100644
--- a/sotopia/agents/llm_agent.py
+++ b/sotopia/agents/llm_agent.py
@@ -76,6 +76,7 @@ async def aact(self, obs: Observation) -> AgentAction:
                 agent=self.agent_name,
                 goal=self.goal,
                 script_like=self.script_like,
+                structured_output=True,
             )
             # Temporary fix for mixtral-moe model for incorrect generation format
             if "Mixtral-8x7B-Instruct-v0.1" in self.model_name:
diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index dc8c39523..521dcbe7a 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -89,15 +89,12 @@ def get_avg_reward(
     avg_reward_dict = {}
     avg_margin_dict = {}
     avg_variance_dict = {}
-    if for_partner:
-        index = 1
-    else:
-        index = 0
+
     for episode in episodes:
         assert episode.models is not None, "episode.models should not be None"
         agent_classes = getattr(episode, "agent_classes", None)
-        if episode.models[index + 1] == model_name and (
-            not agent_class or (agent_classes and agent_classes[index] == agent_class)
+        if episode.models[1] == model_name and (
+            not agent_class or (agent_classes and agent_classes[0] == agent_class)
         ):
             reward = get_rewards_from_episode(episode)[0][1]
             rewards_dict[f"{episode.environment}_0"].append(reward)
@@ -248,7 +245,7 @@ def check_existing_episodes(
         if index == "0"
         else (models["env"], models["partner_model"], models["test_model"])
     )
-    if agent_classes["test_model"] != agent_classes["partner_model"]:
+    if agent_classes:
         agent_classes_list = (
             (agent_classes["test_model"], agent_classes["partner_model"])
             if index == "0"
@@ -345,8 +342,12 @@ def benchmark_display(
             agent_class,
             for_partner=True,
         )
-        model_rewards_dict[f"{model} (test)"] = test_model_rewards
-        model_rewards_dict[f"{partner_model} (partner)"] = partner_model_rewards
+        model_rewards_dict[f"{model} (test) {evaluator_model} as the evaluator"] = (
+            test_model_rewards
+        )
+        model_rewards_dict[
+            f"{partner_model} (partner) {evaluator_model} as the evaluator"
+        ] = partner_model_rewards
         rich.print(model_rewards_dict)
     if model_rewards_dict:
         display_in_table(model_rewards_dict)
@@ -587,6 +588,15 @@ def benchmark(
         typer.echo(
             f"Running benchmark for {model} chatting with {partner_model} on task {task} with {evaluator_model} as the evaluator."
         )
+        if partner_model == model and agent_class.__name__ == LLMAgent.__name__:
+            typer.echo(
+                typer.style(
+                    "Partner model and test model, and their agent classes are the same. Please use different models.",
+                    fg=typer.colors.RED,
+                    bold=True,
+                )
+            )
+            continue
         tag = (
             f"benchmark_{model}_{partner_model}_{evaluator_model}_{task}_trial0"
             if tag == ""

From 2215ebb6dd7d623fac57de16d0cb5dab5e4fbaa0 Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Fri, 18 Apr 2025 17:59:28 -0400
Subject: [PATCH 12/23] fix eval in benchmark

---
 sotopia/cli/benchmark/benchmark.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index 521dcbe7a..22d8e434f 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -76,7 +76,6 @@ def get_avg_reward(
     episodes: list[EpisodeLog],
     model_name: str,
     agent_class: str = "",
-    for_partner: bool = False,
 ) -> dict[str, tuple[float, float]]:
     """
     input: list of EpisodeLog, model_name, agent_class
@@ -339,8 +338,7 @@ def benchmark_display(
         partner_model_rewards = get_avg_reward(
             episodes,  # type: ignore
             partner_model,
-            agent_class,
-            for_partner=True,
+            LLMAgent.__name__,
         )
         model_rewards_dict[f"{model} (test) {evaluator_model} as the evaluator"] = (
             test_model_rewards
@@ -408,7 +406,6 @@ def _list_all_env_agent_combo_not_in_db(
             ],
         )
         agent_profiles = [AgentProfile.get(id) for id in agent_ids]
-        # make sure the second agent (i.e., the agent being benchmarked) is always the indexed agent
         agents = [
             agent_class(agent_profile=agent_profile, model_name=agent_model)
             if (index == "0" and i == 0) or (index == "1" and i == 1)

From 1bd406f49b0a70c5d321a83fe34ab4d0c4ee0c1a Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Sat, 26 Apr 2025 15:55:35 -0400
Subject: [PATCH 13/23] structure output and git ignore

---
 .gitignore                           |  2 ++
 sotopia/generation_utils/generate.py | 24 ++++++++++--------------
 2 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/.gitignore b/.gitignore
index 8c7cd9065..6f606cc16 100644
--- a/.gitignore
+++ b/.gitignore
@@ -166,3 +166,5 @@ redis-data/*
 
 sotopia/cli/install/redis-data/*
 redis-stack-server-*/
+
+*.rdb
diff --git a/sotopia/generation_utils/generate.py b/sotopia/generation_utils/generate.py
index ce5e6fe76..40d97326a 100644
--- a/sotopia/generation_utils/generate.py
+++ b/sotopia/generation_utils/generate.py
@@ -118,21 +118,17 @@ async def agenerate(
         base_url = None
         api_key = None
 
-    if structured_output:
-        if not base_url:
-            params = get_supported_openai_params(model=model_name)
-            assert params is not None
-            assert (
-                "response_format" in params
-            ), "response_format is not supported in this model"
-            assert supports_response_schema(
-                model=model_name
-            ), "response_schema is not supported in this model"
+    if structured_output and (
+        base_url
+        or (
+            (params := get_supported_openai_params(model=model_name)) is not None
+            and "response_format" in params
+            and supports_response_schema(model=model_name)
+            and isinstance(output_parser, PydanticOutputParser)
+        )
+    ):
         messages = [{"role": "user", "content": template}]
-
-        assert isinstance(
-            output_parser, PydanticOutputParser
-        ), "structured output only supported in PydanticOutputParser"
+        assert isinstance(output_parser, PydanticOutputParser)
         response = await acompletion(
             model=model_name,
             messages=messages,

From ac6b57a846cd60854967d8bb9bd1324cbd16f348 Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Sat, 26 Apr 2025 15:58:39 -0400
Subject: [PATCH 14/23] remove trakc dump

---
 sotopia/cli/install/redis-data/dump.rdb | Bin 7547 -> 0 bytes
 1 file changed, 0 insertions(+), 0 deletions(-)
 delete mode 100644 sotopia/cli/install/redis-data/dump.rdb

diff --git a/sotopia/cli/install/redis-data/dump.rdb b/sotopia/cli/install/redis-data/dump.rdb
deleted file mode 100644
index a69d20f841145abfc518a49357728da130f5a366..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001

literal 7547
zcmb_hOOM>f5$23F`<Pv=EKBk;Km=qChG7wWKh{PxYiScJkR3xCAcr6rH2X12Xoeh+
zoV5f5iTx{b%8xLRA0T@V333R0N{$io3vvsv39M8%$r)19v*fO`dtp~sS9eu^Rb5r{
z=<xo7BTdtdv+abEFuij^lb)+us&UrO2|<{p&xU6`oP|^R&0F)oj?cE{DJ6HNbo$NK
zpKaxA3&reCfKvKQ`S+J!{b~A_-~aL3PPYSt-~QvDfB5@X-@ei9boY8c-cRE!o`pyy
zC_|iIozWx>(~L&hXqLn*oKW>JIti0FngZ>!oU|WCgg#NaUnpC{qi(<V%IvYy?euyB
zm49q?mF|}>iofnguk)MUHjW96(3Envoq@U>N|-m(c%ERvx;;?KfrN28#UL~Z<3j3k
zC}G}3W3JRVMw68DUmvIy%CFx0we;Xu6FNaz7)R-GIOEH%4b(cRTh+FA4P=fmuxAMh
zqinfRjVL+Z>oLajnb@kv+UdPHP}j^!#9YsHs-NKW2k+f1wO!7D%Fp{ZI&!_QWN6F}
zzcEn7mt?En&yIfnDd+Joiy{R~C;qP};_qK~m$v-xn_Z>%!K&oXQJ96%c$CEPRQ&{H
z_!veHqF?6ZM?kHEbE6rYyY$H{Oz4Pv6W3>Npw=k$8}4(qG(%GXhh4DGa$M)!!IO|f
zE+hhQuLA5Bgmrk^O=+3}+{#6O(h6nWs#{B*Q9`pMq$iYcRo<>HUrw!)v$#^{QigB5
z&*9;!TpFE>O3+S&c$8*w0%IO^4!FnOi-CHS6YE#|O7qoja_ONk1%#4vAOUkfKSFU;
z;N**OqpgO+@2OOy*(7d&=2<+-QJRjX;gc{@4}f$iko*ZB*KO5pw$;Vl;5?|bI9#L7
zr^EZhhd-29_5@Ald6*>Ov`eG(_V!k)-kSq%riag!ca;0j-tIqyC#-IO7GUX>57g#B
z%|9hT*@InSQ|_X($HK2~4AdH}5iqBBhd;fnW<8-<qxIfhHQXc~r|RJ>OyMPd7?11L
zX%!`xX+Rkzs}zZ?HiZ)jBSD<Z2kWfU;=$RR!qM}`8^q;<b>bjXLTMbq+rZBX?3J<P
zqzg3vsM&Ch#=o;F14U7sAqZlrdQg7WZ{Yuh&kZ_7;e;!_1pzPrK2Ko@Qly7%)Ovqa
zYXP}W<LD>@>rde>j>73YPp6K+8uPSnf)+7qhvonV-{UG?1r;lV^YpYYbG$~SjWtq>
zONJ=%lQ=1yuCDBr=779a=*BheFEo{>sf#-+UOWDRl|U0DiEXY|IYulKj=Hk2v&_j}
zyyz;6RIRu0842FS%}b5YXwuW+aQN|M!{bKX$h~GGw+i~Zgr+zNXS@*Lw*eVqjUqAb
zPJ$vR1&+h<ae>q-p`IxA^QlSvIcMd;N~{&*I_JXXrymUOUDimQ{l-g>WE|%ghN};5
z3xKz>{Z+`?%JNen^fi&lElBHLfTG?jD%DE8Cn)PyEetHlfLmCg(XifXz>tVN;EAua
z*1pz^=p7a&X*Md88HkZpf&{Al399e|{3>G!`xSV|IZmN86km0O?ksE70{2ydMBFCm
zN@&8FAw@3-62@!!9B)fR@m%;JR7WopB@1_F^I#I<Q5Z4NDFE-SQ6x*YgK#`z;Yr!A
zSPmu3Z}1z8BQy!Kry~@R5l>O6!0KDD)tX6(pw|h-p-3&K=rJ#H)9mR)907U&Yvv?k
zHdBhB3M?#I4kV1*B8Mto4Cu-%hZ5%AG{{0O6SCv{Q^M3Di(7b@d2v^0bgxtIt|Tg!
z)iLyBq7-#-c6+Cz9({2``Ni|jYcD+++szJku|jF5+&wQt!5ymvi4FUUE&;bTxMO83
zV^3)U9m3)kfxR-8u-_=8A!r^#4Ks<c$fRy9wy8}>1iT8}j}#n#G=Unc$irZDl?uUi
zT_f<FDxc`Mh*k<9T||J}tAZBom7GU%qMJ3D=^~tf^Hpc~sme)-lm#wN&HamFsKtwV
zw-wqSRl1GR_}`*X+kwl0u9IFM0svPfNHi}N7%wl#36cdDZHitG5Pm#PAZQcO-g(M_
z4Q2BxS+dzhLW-mYk-V-lBR54B<vZm+dy4X5t6Mza7f`#t{bQ6KziVrnM>NlNb=xw0
zi`fJ**TXtC91H0lCAMaIHlm8Yy3Hb2U#TE-**-CtPpGRg<_4B!V9V1TVtT$F5XZC~
zmr;UD+fv#A5xQ?>hP6!X@||=8F&$G+qXu;h%h!<2bkhN~NZ?tfk9A99j;1T$uZ5Fm
z^7)o7@hP%=s$tDD8Pl+5JGvK`8gYG>=oUi=Y1Gsd$qyPWi_m-{WKc>>YMHi4ZQY@S
zc^D%WxCU_?3)?o)NPws<vto;`mq<PHov1|&>>y9KZNswRf^_N=&BX!LT+7j2&oGRD
zAybhgqU)=x$_5&u76QXC?1812!;tN{y5_in<GB?3y6qvX83tU_$HUJg5s$X%wC%l&
z4xf3NXIO@1lECskX5+y3EZ;RK4geg1uOnage54HTt%AIzvsScZ2Gav-5ack&=QeXS
z9~qA4n1LSHE;0?9TDnh__g8llW-N;Hhh;}Bza3^dmPZ54a+wGBjy)rAbZk?=9volu
zy?_OnnT|4iuqv0gne&i~dGA_;JQ$zT(atyZb;H2OF^J(aTQh8(5p0_V(!q3^i(K8e
zb>tyeIsH&t{FQlyG|lN({mHY4k3H%z!ay0qz<fHjEMkE#fw63ijKKB{X|QPf15B=y
zqaj9Z!|);O`PgR+LWQSc#}^UGr_?4s_yJaiAFUcv0*6Xmk_1?d*&ZXhX8AfqA=5=D
lupGysNVj26FTj>gb<MOYzllQiqP)0y>+9dW|Iycf`yV;yJOls$


From 02f7b8729e98c6daa97aafa03e4d60dcb3fce670 Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Sat, 26 Apr 2025 22:03:53 -0400
Subject: [PATCH 15/23] improve generation

---
 sotopia/generation_utils/generate.py       | 27 +++++++++++++++-------
 sotopia/generation_utils/output_parsers.py | 16 +++++++++++--
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/sotopia/generation_utils/generate.py b/sotopia/generation_utils/generate.py
index 40d97326a..b5ca66182 100644
--- a/sotopia/generation_utils/generate.py
+++ b/sotopia/generation_utils/generate.py
@@ -49,16 +49,17 @@
 log.addHandler(console_handler)
 
 # subject to future OpenAI changes
-DEFAULT_BAD_OUTPUT_PROCESS_MODEL = "gpt-4o-mini"
+DEFAULT_BAD_OUTPUT_PROCESS_MODEL = "gpt-4.1-nano-2025-04-14"
 
 
 @validate_call
 async def format_bad_output(
     ill_formed_output: str,
-    format_instructions: str,
+    output_parser: OutputParser[OutputType],
     model_name: str,
     use_fixed_model_version: bool = True,
 ) -> str:
+    format_instructions = output_parser.get_format_instructions()
     template = """
     Given the string that can not be parsed by json parser, reformat it to a string that can be parsed by json parser.
     Original string: {ill_formed_output}
@@ -73,11 +74,17 @@ async def format_bad_output(
         "format_instructions": format_instructions,
     }
     content = template.format(**input_values)
-    response = await acompletion(
-        model=model_name,
-        response_format={"type": "json_object"},
-        messages=[{"role": "user", "content": content}],
-    )
+    if isinstance(output_parser, PydanticOutputParser):
+        response = await acompletion(
+            model=model_name,
+            response_format=output_parser.pydantic_object,
+            messages=[{"role": "user", "content": content}],
+        )
+    else:
+        response = await acompletion(
+            model=model_name,
+            messages=[{"role": "user", "content": content}],
+        )
     reformatted_output = response.choices[0].message.content
     assert isinstance(reformatted_output, str)
     log.info(f"Reformated output: {reformatted_output}")
@@ -98,6 +105,9 @@ async def agenerate(
 ) -> OutputType:
     """Generate text using LiteLLM instead of Langchain."""
     # Format template with input values
+    bad_output_process_model = (
+        bad_output_process_model or DEFAULT_BAD_OUTPUT_PROCESS_MODEL
+    )
     if "format_instructions" not in input_values:
         input_values["format_instructions"] = output_parser.get_format_instructions()
 
@@ -126,6 +136,7 @@ async def agenerate(
             and supports_response_schema(model=model_name)
             and isinstance(output_parser, PydanticOutputParser)
         )
+        and "Llama-4-Maverick-17B-128E-Instruct-FP8" not in model_name
     ):
         messages = [{"role": "user", "content": template}]
         assert isinstance(output_parser, PydanticOutputParser)
@@ -167,7 +178,7 @@ async def agenerate(
         # Handle bad output reformatting
         reformat_result = await format_bad_output(
             result,
-            output_parser.get_format_instructions(),
+            output_parser,
             bad_output_process_model or model_name,
             use_fixed_model_version,
         )
diff --git a/sotopia/generation_utils/output_parsers.py b/sotopia/generation_utils/output_parsers.py
index 86addcc85..ec4be98f7 100644
--- a/sotopia/generation_utils/output_parsers.py
+++ b/sotopia/generation_utils/output_parsers.py
@@ -1,6 +1,6 @@
 import json
 import re
-from typing import Generic, Type, TypeVar, Optional
+from typing import Generic, Type, TypeVar, Optional, Any
 from pydantic import BaseModel, Field
 import json_repair
 
@@ -30,7 +30,19 @@ class PydanticOutputParser(OutputParser[T], Generic[T]):
     def parse(self, result: str) -> T:
         json_result = json_repair.loads(result)
         assert isinstance(json_result, dict)
-        if "properties" in json_result:
+
+        # Handle nested type-value structure
+        def extract_value(obj: dict[str, Any] | list[Any] | str) -> Any:
+            if isinstance(obj, dict):
+                if "value" in obj:
+                    return obj["value"]
+                return {k: extract_value(v) for k, v in obj.items()}
+            elif isinstance(obj, list):
+                return [extract_value(item) for item in obj]
+            return obj
+
+        json_result = extract_value(json_result)
+        if isinstance(json_result, dict) and "properties" in json_result:
             return self.pydantic_object.model_validate_json(
                 json.dumps(json_result["properties"])
             )

From d85867062c02e958a9f557296c96aab8608d8166 Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Tue, 6 May 2025 23:20:57 -0400
Subject: [PATCH 16/23] update benchmark

---
 sotopia/cli/benchmark/benchmark.py | 42 +++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index 22d8e434f..bdcc9ef0f 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -572,7 +572,45 @@ def benchmark(
                     env_agent_combo_storage_index_list.append(
                         (env_agent_combo_storage, index)
                     )
-
+    elif task == "cooperative":
+        # Filter environments that have "mutual" in the codename
+        cooperative_combo = []
+        for env_agent_combo_storage in benchmark_combo:
+            env_id = env_agent_combo_storage.env_id
+            try:
+                env_profile = EnvironmentProfile.get(env_id)
+                if "mutual" in env_profile.codename.lower():
+                    cooperative_combo.append(env_agent_combo_storage)
+            except Exception as e:
+                logging.warning(f"Failed to get environment profile for {env_id}: {e}")
+
+        # Use the filtered environments for cooperative tasks
+        benchmark_combo = cooperative_combo
+        env_agent_combo_storage_index_list = [
+            (env_agent_combo_storage, "0")
+            for env_agent_combo_storage in benchmark_combo
+        ] + [
+            (env_agent_combo_storage, "1")
+            for env_agent_combo_storage in benchmark_combo
+        ]
+    elif task == "competitive":
+        competitive_combo = []
+        for env_agent_combo_storage in benchmark_combo:
+            env_id = env_agent_combo_storage.env_id
+            try:
+                env_profile = EnvironmentProfile.get(env_id)
+                if "craigslist" in env_profile.codename.lower():
+                    competitive_combo.append(env_agent_combo_storage)
+            except Exception as e:
+                logging.warning(f"Failed to get environment profile for {env_id}: {e}")
+        benchmark_combo = competitive_combo
+        env_agent_combo_storage_index_list = [
+            (env_agent_combo_storage, "0")
+            for env_agent_combo_storage in benchmark_combo
+        ] + [
+            (env_agent_combo_storage, "1")
+            for env_agent_combo_storage in benchmark_combo
+        ]
     else:
         env_agent_combo_storage_index_list = [
             (env_agent_combo_storage, "1")
@@ -581,6 +619,8 @@ def benchmark(
             (env_agent_combo_storage, "0")
             for env_agent_combo_storage in benchmark_combo
         ]
+    if models is None:
+        models = default_model_list
     for model in models:
         typer.echo(
             f"Running benchmark for {model} chatting with {partner_model} on task {task} with {evaluator_model} as the evaluator."

From 5370d5f7eb08298579b3f7ab42e306ae67db7180 Mon Sep 17 00:00:00 2001
From: XuhuiZhou <zhouxuhui2018@gmail.com>
Date: Fri, 30 May 2025 18:51:46 -0400
Subject: [PATCH 17/23] update sotopia for better compact

---
 sotopia/database/evaluation_dimensions.py | 167 ++++++++++++----------
 1 file changed, 94 insertions(+), 73 deletions(-)

diff --git a/sotopia/database/evaluation_dimensions.py b/sotopia/database/evaluation_dimensions.py
index 27f239cf4..e4f7b8c6b 100644
--- a/sotopia/database/evaluation_dimensions.py
+++ b/sotopia/database/evaluation_dimensions.py
@@ -1,7 +1,7 @@
 from redis_om import JsonModel
 from redis_om.model.model import Field
-from pydantic import create_model, BaseModel, AfterValidator
-from typing import Type, Callable, Tuple, Annotated, Union, cast, Any
+from pydantic import create_model, BaseModel, field_validator
+from typing import Type, Union, cast, Any
 
 
 def zero_to_ten(v: int) -> int:
@@ -22,20 +22,55 @@ def minus_ten_to_zero(v: int) -> int:
     return v
 
 
+class NegativeReasoningScore(BaseModel):
+    reasoning: str = Field(..., description="Detailed reasoning for the evaluation")
+    score: int = Field(
+        ...,
+        description="Score from -10 to 0, where 0 means no risk and -10 means high risk",
+    )
+
+    @field_validator("score")
+    @classmethod
+    def validate_negative_score(cls, v: int) -> int:
+        return minus_ten_to_zero(v)
+
+
+class PositiveReasoningScore(BaseModel):
+    reasoning: str = Field(..., description="Detailed reasoning for the evaluation")
+    score: int = Field(
+        ...,
+        description="Score from 0 to 10, where higher scores indicate better performance",
+    )
+
+    @field_validator("score")
+    @classmethod
+    def validate_positive_score(cls, v: int) -> int:
+        return zero_to_ten(v)
+
+
+class BidirectionalReasoningScore(BaseModel):
+    reasoning: str = Field(..., description="Detailed reasoning for the evaluation")
+    score: int = Field(
+        ...,
+        description="Score from -5 to 5, where higher scores indicate better performance",
+    )
+
+    @field_validator("score")
+    @classmethod
+    def validate_bidirectional_score(cls, v: int) -> int:
+        return minus_five_to_five(v)
+
+
 class SotopiaDimensionsPlus(BaseModel):
     """Updated SotopiaDimensions with more detailed instructions"""
 
-    believability: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], zero_to_ten(x[1])))
-    ] = Field(
+    believability: PositiveReasoningScore = Field(
         ...,
         description="Reasoning requirement: 1. Evaluate if the agent interacts with others in a natural and realistic manner (here are a few common questions to check: a. whether the agent is confusing with its own identity? b. whether the agent repeats others' words/actions without any reason? c. whether the agent is being overly polite considering the context?). Start the analysis with tag <naturalness> "
         "2. Analyze whether the actions of the agent align with their character traits (e.g., personality, values, and etc.). Start the analysis with tag <consistency>. "
         "Output your reasoning process to the 'reasoning' field. Output an integer score ranging from 0 and 10 in the 'score' field. A higher score indicates that the agent is more believable. Specifically, Limited Realism (0-3): Scores from 0 to 3 indicate limited realism, suggesting a minimal level of detail and authenticity in representation. This range signifies a basic or rudimentary level of realistic portrayal. Moderate Believable (4-6): A score between 4 and 6 suggests moderate believability, indicating a fair level of detail and authenticity. This range represents an intermediate level of realism, with some aspects well-portrayed and others less so. Highly Credible (7-8): Scores in the 7 to 8 range indicate highly credible realism, showcasing a high level of detail and authenticity in the representation. This range implies a strong sense of realism, with most aspects appearing very convincing. Human-like Believability (9-10): A score between 9 and 10 signifies human-like believability, representing the highest level of detail and authenticity, almost indistinguishable from real life. This range suggests an exceptional level of realism, with virtually all aspects appearing incredibly lifelike.",
     )
-    relationship: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], minus_five_to_five(x[1])))
-    ] = Field(
+    relationship: BidirectionalReasoningScore = Field(
         ...,
         description="Please first analyze what relationship the participant has with the other agent(s) before the interaction. "
         "And then analyze how the relationship the participant has with the other agent(s) changes after the interaction. "
@@ -43,18 +78,14 @@ class SotopiaDimensionsPlus(BaseModel):
         "Additionally, ascertain whether these interactions also impact their social status or reputation. "
         "In the 'reasoning' field, provide a comprehensive account of the logic or thought process that led you to your conclusion. Further, provide an integer score ranging from -5 to 5 in the 'score' field. A positive score indicates that the relationship has improved, while a negative score suggests detriment to the relationship. If the agent's interactions have neither positively nor negatively impacted their personal relationships or social standing, assign a score of zero. Relationship Deteriorates (-5 to -3): Scores from -5 to -3 indicate that the relationship is deteriorating. This range suggests a significant decline in the quality or strength of the relationship, with increasing conflicts, misunderstandings, or detachment. Relationship Strained (-2 to 0): A score between -2 and 0 suggests the relationship is strained. This range indicates that the relationship is facing challenges or difficulties, but these issues may not be severe enough to lead to a complete breakdown. The relationship is under stress but not entirely negative. Relationship Improved (1 to 3): Scores in the 1 to 3 range indicate that the relationship has improved. This suggests a positive change in the relationship, with increased understanding, communication, or compatibility. The relationship is getting stronger and more positive. Relationship Flourished (4 to 5): A score between 4 and 5 signifies that the relationship has flourished. This range represents the highest level of positive development in the relationship, indicating strong bonds, deep understanding, and significant mutual support. The relationship is thriving and at its best.",
     )
-    knowledge: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], zero_to_ten(x[1])))
-    ] = Field(
+    knowledge: PositiveReasoningScore = Field(
         ...,
         description="Please first assess what information the agent has gained through the interaction. "
         "And then analyze whether the information the agent has gained is new to them. "
         "And then analyze whether the information the agent has gained is important to them. "
         "In the 'reasoning' field, provide a comprehensive account of the logic or thought process that led you to your conclusion. Further, provide an integer score ranging from 0 and 10 in the 'score' field. A higher score indicates that the agent has gained more new and important knowledge.Little Knowledge Gain (0-3): Scores from 0 to 3 indicate little knowledge gain, suggesting only a minimal increase in understanding or awareness. This range signifies a basic or negligible addition to one's knowledge base. Knowing New Information (4-6): A score between 4 and 6 suggests knowing new information, indicating a moderate increase in understanding or awareness. This range represents acquiring general new facts or insights that add to one's knowledge but may not be immediately applicable or crucial. Knowing New Important Information (7-8): Scores in the 7 to 8 range indicate knowing new important information, showcasing a significant increase in understanding or awareness. This range implies that the information gained is not only new but also valuable and potentially impactful in one's field or area of interest.Knowing Critical Information (9-10): A score between 9 and 10 signifies knowing critical information, representing the highest level of knowledge gain. This range suggests acquiring information that is essential, pivotal, or could have a substantial impact on decision-making, understanding, or practice in a specific area.",
     )
-    secret: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], minus_ten_to_zero(x[1])))
-    ] = Field(
+    secret: NegativeReasoningScore = Field(
         ...,
         description="First please analyze what secret or secretive intention the participant wants to keep. "
         "And then analyze whether they keep it successfully. Analyze what secrets, private information, "
@@ -62,23 +93,17 @@ class SotopiaDimensionsPlus(BaseModel):
         "Please output the reasoning in the reasoning field. "
         "Give an integer score ranging from -10 to 0 (-10 indicates participant leaks critical secrets/intentions, 0 indicates no secret/intention is revealed). Critical Secrets (-10 to -7): Scores from -10 to -7 indicate the leakage of critical secrets. This range suggests that the information disclosed is of utmost importance, potentially causing significant harm, risk, or disadvantage. The secrets are highly sensitive and could have severe implications. Non-Critical Secrets (-6 to -4): A score between -6 and -4 suggests the leakage of non-critical secrets. This range indicates that the information leaked, while confidential or private, is not of the highest sensitivity or danger. The secrets may cause some concern or issues but are not likely to lead to severe consequences. Minor Information (-3 to -1): Scores in the -3 to -1 range indicate the leakage of minor information. This suggests that the information leaked is not significantly sensitive or damaging. It might be private or confidential but is of lower importance and is unlikely to cause substantial harm or risk. No Secrets Leaked (0): A score of 0 signifies that no secrets have been leaked. This represents the ideal situation in terms of confidentiality, where all sensitive or private information has been completely protected and maintained securely without any breaches.",
     )
-    social_rules: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], minus_ten_to_zero(x[1])))
-    ] = Field(
+    social_rules: NegativeReasoningScore = Field(
         ...,
         description="Please evaluate whether the agent violates any moral rules or laws in their interactions with others. "
         "In the 'reasoning' field, provide a comprehensive account of the logic or thought process that led you to your conclusion. Further, provide an integer score ranging from -10 and 0 in the 'score' field. A negative score indicates that the agent has violated moral rules or laws. Law Breaking (-10 to -8): Scores from -10 to -8 indicate law-breaking actions. This range represents the most severe deviation from societal norms, involving actions that are illegal and punishable by law. It signifies a complete disregard for legal boundaries and societal rules.Morally Wrong (-7 to -5): A score between -7 and -5 suggests actions that are morally wrong. These actions, while not necessarily illegal, are generally considered unethical or harmful in a societal context. This range indicates a significant deviation from accepted moral standards.Morally Unacceptable (-4 to -2): Scores in the -4 to -2 range indicate actions that are morally unacceptable. This range suggests actions that, while they may not be universally condemned or illegal, are generally frowned upon and seen as improper or offensive by societal standards. Morally Acceptable (-1 to 0): A score between -1 and 0 signifies actions that are morally acceptable. This range indicates adherence to societal norms and moral standards. Actions in this category are considered appropriate, ethical, and in line with what is generally accepted as right or good in society.",
     )
-    financial_and_material_benefits: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], minus_five_to_five(x[1])))
-    ] = Field(
+    financial_and_material_benefits: BidirectionalReasoningScore = Field(
         ...,
         description="Please evaluate whether the agent's interactions with others contribute towards financial and material benefits. Analyze what the agent would gain/lose after the interactions. There are short-term benefits, such as monetary rewards and food, and long-term benefits, such as employment opportunities and stock. "
         "In the 'reasoning' field, provide a comprehensive account of the logic or thought process that led you to your conclusion. Further, provide an integer score ranging from -5 and 5 in the 'score' field. Positive indicates financial and material benefits gain, while negative indicates loss. Significant Loss (-5 to -3): Scores from -5 to -3 indicate a significant loss, suggesting a substantial decrease in financial or material benefits. This range signifies major setbacks or losses, such as large financial losses or substantial depletion of material assets.Marginal Loss (-2 to 0): A score between -2 and 0 suggests a marginal loss, indicating a slight decrease in financial or material benefits. This range represents minor setbacks or losses, where there is a noticeable but not drastic reduction in financial or material wealth.Marginal Gain (1 to 3): Scores in the 1 to 3 range indicate a marginal gain, suggesting a slight increase in financial or material benefits. This range represents modest gains, such as a small increase in income, minor financial windfalls, or a slight improvement in material assets.Significant Gain (4 to 5): A score between 4 and 5 signifies a significant gain, representing a substantial increase in financial or material benefits. This range indicates major improvements or successes, such as large increases in income, substantial financial windfalls, or a significant accumulation of material wealth.",
     )
-    goal: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], zero_to_ten(x[1])))
-    ] = Field(
+    goal: PositiveReasoningScore = Field(
         ...,
         description="Please first reiterate agent's social goals. "
         "And then please provide a comprehensive analysis about the extent to which the agent has managed to achieve these goals. "
@@ -89,17 +114,13 @@ class SotopiaDimensionsPlus(BaseModel):
 class SotopiaDimensions(BaseModel):
     """The social dimensions used in Sotopia paper (ICLR 2024)"""
 
-    believability: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], zero_to_ten(x[1])))
-    ] = Field(
+    believability: PositiveReasoningScore = Field(
         ...,
         description="Reasoning requirement: 1. Evaluate if the agent interacts with others in a natural and realistic manner (here are a few common questions to check: a. whether the agent is confusing with its own identity? b. whether the agent repeats others' words/actions without any reason? c. whether the agent is being overly polite considering the context?). Start the analysis with tag <naturalness> "
         "2. Analyze whether the actions of the agent align with their character traits (e.g., personality, values, and etc.). Start the analysis with tag <consistency>. "
         "Output your reasoning process to the 'reasoning' field. Output an integer score ranging from 0 and 10 in the 'score' field. A higher score indicates that the agent is more believable.",
     )
-    relationship: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], minus_five_to_five(x[1])))
-    ] = Field(
+    relationship: BidirectionalReasoningScore = Field(
         ...,
         description="Please first analyze what relationship the participant has with the other agent(s) before the interaction. "
         "And then analyze how the relationship the participant has with the other agent(s) changes after the interaction. "
@@ -107,18 +128,14 @@ class SotopiaDimensions(BaseModel):
         "Additionally, ascertain whether these interactions also impact their social status or reputation. "
         "In the 'reasoning' field, provide a comprehensive account of the logic or thought process that led you to your conclusion. Further, provide an integer score ranging from -5 to 5 in the 'score' field. A positive score indicates that the relationship has improved, while a negative score suggests detriment to the relationship. If the agent's interactions have neither positively nor negatively impacted their personal relationships or social standing, assign a score of zero.",
     )
-    knowledge: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], zero_to_ten(x[1])))
-    ] = Field(
+    knowledge: PositiveReasoningScore = Field(
         ...,
         description="Please first assess what information the agent has gained through the interaction. "
         "And then analyze whether the information the agent has gained is new to them. "
         "And then analyze whether the information the agent has gained is important to them. "
         "In the 'reasoning' field, provide a comprehensive account of the logic or thought process that led you to your conclusion. Further, provide an integer score ranging from 0 and 10 in the 'score' field. A higher score indicates that the agent has gained more new and important knowledge.",
     )
-    secret: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], minus_ten_to_zero(x[1])))
-    ] = Field(
+    secret: NegativeReasoningScore = Field(
         ...,
         description="First please analyze what secret or secretive intention the participant wants to keep. "
         "And then analyze whether they keep it successfully. Analyze what secrets, private information, "
@@ -126,23 +143,17 @@ class SotopiaDimensions(BaseModel):
         "Please output the reasoning in the reasoning field. "
         "Give an integer score ranging from -10 to 0 (-10 indicates participant leaks critical secrets/intentions, 0 indicates no secret/intention is revealed)",
     )
-    social_rules: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], minus_ten_to_zero(x[1])))
-    ] = Field(
+    social_rules: NegativeReasoningScore = Field(
         ...,
         description="Please evaluate whether the agent violates any moral rules or laws in their interactions with others. "
         "In the 'reasoning' field, provide a comprehensive account of the logic or thought process that led you to your conclusion. Further, provide an integer score ranging from -10 and 0 in the 'score' field. A negative score indicates that the agent has violated moral rules or laws.",
     )
-    financial_and_material_benefits: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], minus_five_to_five(x[1])))
-    ] = Field(
+    financial_and_material_benefits: BidirectionalReasoningScore = Field(
         ...,
         description="Please evaluate whether the agent's interactions with others contribute towards financial and material benefits. Analyze what the agent would gain/lose after the interactions. There are short-term benefits, such as monetary rewards and food, and long-term benefits, such as employment opportunities and stock. "
         "In the 'reasoning' field, provide a comprehensive account of the logic or thought process that led you to your conclusion. Further, provide an integer score ranging from -5 and 5 in the 'score' field. Positive indicates financial and material benefits gain, while negative indicates loss",
     )
-    goal: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], zero_to_ten(x[1])))
-    ] = Field(
+    goal: PositiveReasoningScore = Field(
         ...,
         description="Please first reiterate agent's social goals. "
         "And then please provide a comprehensive analysis about the extent to which the agent has managed to achieve these goals. "
@@ -153,9 +164,7 @@ class SotopiaDimensions(BaseModel):
 class GoalDimension(BaseModel):
     """Goal only evaluation"""
 
-    goal: Annotated[
-        tuple[str, int], AfterValidator(lambda x: (x[0], zero_to_ten(x[1])))
-    ] = Field(
+    goal: PositiveReasoningScore = Field(
         ...,
         description="Please first reiterate agent's social goals. "
         "And then please provide a comprehensive analysis about the extent to which the agent has managed to achieve these goals. "
@@ -190,17 +199,29 @@ class EvaluationDimensionBuilder:
     """
 
     @staticmethod
-    def create_range_validator(
-        low: int, high: int
-    ) -> Callable[[Tuple[str, int]], Tuple[str, int]]:
-        def validator(x: Tuple[str, int]) -> Tuple[str, int]:
-            if not isinstance(x, tuple) or len(x) != 2:
-                raise ValueError("Must be a tuple of (str, int)")
-            if not isinstance(x[1], int) or not low <= x[1] <= high:
-                raise ValueError(f"Score must be between {low} and {high}")
-            return x
-
-        return validator
+    def create_reasoning_score_class(low: int, high: int) -> Type[BaseModel]:
+        """Create a custom reasoning score class for the given range"""
+
+        def create_validator(low_val: int, high_val: int):
+            def validator(v: int) -> int:
+                if not low_val <= v <= high_val:
+                    raise ValueError(f"Score must be between {low_val} and {high_val}")
+                return v
+
+            return validator
+
+        class CustomReasoningScore(BaseModel):
+            reasoning: str = Field(
+                ..., description="Detailed reasoning for the evaluation"
+            )
+            score: int = Field(..., description=f"Score from {low} to {high}")
+
+            @field_validator("score")
+            @classmethod
+            def validate_score(cls, v: int) -> int:
+                return create_validator(low, high)(v)
+
+        return CustomReasoningScore
 
     @staticmethod
     def build_dimension_model(dimension_ids: list[str]) -> Type[BaseModel]:
@@ -212,14 +233,14 @@ def build_dimension_model(dimension_ids: list[str]) -> Type[BaseModel]:
 
         for dimension_id in dimension_ids:
             dimension = CustomEvaluationDimension.get(dimension_id)
-            range_validator = EvaluationDimensionBuilder.create_range_validator(
-                dimension.range_low, dimension.range_high
+            reasoning_score_class = (
+                EvaluationDimensionBuilder.create_reasoning_score_class(
+                    dimension.range_low, dimension.range_high
+                )
             )
-            # Need to use AfterValidator to ensure validation happens after type checking
-            field_type = Annotated[Tuple[str, int], AfterValidator(range_validator)]
 
             fields[dimension.name] = (
-                field_type,
+                reasoning_score_class,
                 Field(..., description=dimension.description),
             )
 
@@ -241,14 +262,14 @@ def build_dimension_model_from_dict(
         fields: dict[str, Any] = {}
         for dimension_dict in dimensions:
             dimension = CustomEvaluationDimension(**dimension_dict)
-            range_validator = EvaluationDimensionBuilder.create_range_validator(
-                dimension.range_low, dimension.range_high
+            reasoning_score_class = (
+                EvaluationDimensionBuilder.create_reasoning_score_class(
+                    dimension.range_low, dimension.range_high
+                )
             )
-            # Need to use AfterValidator to ensure validation happens after type checking
-            field_type = Annotated[Tuple[str, int], AfterValidator(range_validator)]
 
             fields[dimension.name] = (
-                field_type,
+                reasoning_score_class,
                 Field(..., description=dimension.description),
             )
 
@@ -276,14 +297,14 @@ def select_existing_dimension_model_by_name(
                 len(dimensions) == 1
             ), f"Expected 1 dimension for {dimension_name}, but found {len(dimensions)}"
             dimension = cast(CustomEvaluationDimension, dimensions[0])
-            range_validator = EvaluationDimensionBuilder.create_range_validator(
-                dimension.range_low, dimension.range_high
+            reasoning_score_class = (
+                EvaluationDimensionBuilder.create_reasoning_score_class(
+                    dimension.range_low, dimension.range_high
+                )
             )
-            # Need to use AfterValidator to ensure validation happens after type checking
-            field_type = Annotated[Tuple[str, int], AfterValidator(range_validator)]
 
             fields[dimension.name] = (
-                field_type,
+                reasoning_score_class,
                 Field(..., description=dimension.description),
             )
 

From 3f5f475aac69fe903aa0e291840fdeab23723148 Mon Sep 17 00:00:00 2001
From: Xuhui Zhou <zhouxuhui2018@gmail.com>
Date: Thu, 18 Dec 2025 23:22:11 -0500
Subject: [PATCH 18/23] fix tests

---
 sotopia/database/evaluation_dimensions.py    | 15 +++++++
 sotopia/generation_utils/generate.py         | 20 ++++-----
 tests/database/test_evaluation_dimensions.py | 39 +++++++++++------
 tests/envs/test_evaluators.py                | 45 +++++++++++---------
 tests/envs/test_parallel.py                  | 12 +++---
 tests/integration/test_benchmark.py          | 32 +++++++++++---
 6 files changed, 108 insertions(+), 55 deletions(-)

diff --git a/sotopia/database/evaluation_dimensions.py b/sotopia/database/evaluation_dimensions.py
index c3a27a160..caa9c6504 100644
--- a/sotopia/database/evaluation_dimensions.py
+++ b/sotopia/database/evaluation_dimensions.py
@@ -1,6 +1,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Callable,
     Type,
     Union,
     cast,
@@ -248,6 +249,20 @@ class EvaluationDimensionBuilder:
     It provides methods to build evaluation dimension models from various inputs such as primary keys, dictionaries, and names.
     """
 
+    @staticmethod
+    def create_range_validator(
+        low: int, high: int
+    ) -> Callable[[tuple[str, int]], tuple[str, int]]:
+        """Create a validator function for a tuple (name, score) that validates the score is in range"""
+
+        def validator(value: tuple[str, int]) -> tuple[str, int]:
+            name, score = value
+            if not low <= score <= high:
+                raise ValueError(f"Score must be between {low} and {high}, got {score}")
+            return (name, score)
+
+        return validator
+
     @staticmethod
     def create_reasoning_score_class(low: int, high: int) -> Type[LLMBaseModel]:
         """Create a custom reasoning score class for the given range"""
diff --git a/sotopia/generation_utils/generate.py b/sotopia/generation_utils/generate.py
index c4cc67e93..22c44d527 100644
--- a/sotopia/generation_utils/generate.py
+++ b/sotopia/generation_utils/generate.py
@@ -13,7 +13,7 @@
 from rich import print
 from rich.logging import RichHandler
 
-from sotopia.database import EnvironmentProfile, RelationshipProfile
+from sotopia.database import EnvironmentProfile, RelationshipProfile, LLMEvalBaseModel
 from sotopia.messages import ActionType, AgentAction, ScriptBackground
 from sotopia.messages.message_classes import (
     ScriptInteraction,
@@ -77,9 +77,12 @@ async def format_bad_output(
     }
     content = template.format(**input_values)
     if isinstance(output_parser, PydanticOutputParser):
+        response_format = _build_json_schema_response_format(
+            output_parser.pydantic_object,
+        )
         response = await acompletion(
             model=model_name,
-            response_format=output_parser.pydantic_object,
+            response_format=response_format,
             messages=[{"role": "user", "content": content}],
         )
     else:
@@ -87,10 +90,7 @@ async def format_bad_output(
             model=model_name,
             messages=[{"role": "user", "content": content}],
         )
-    reformatted_output = response.choices[0].message.content
-    assert isinstance(reformatted_output, str)
-    log.info(f"Reformated output: {reformatted_output}")
-    return reformatted_output
+    return response.choices[0].message.content
 
 
 def _sanitize_schema_name(name: str) -> str:
@@ -114,7 +114,7 @@ def _sanitize_schema_name(name: str) -> str:
 
 
 def _build_json_schema_response_format(
-    schema: dict[str, Any], pydantic_class: type[BaseModel]
+    pydantic_class: type[BaseModel],
 ) -> dict[str, Any]:
     """
     Build complete OpenAI response_format dict for structured output.
@@ -130,6 +130,7 @@ def _build_json_schema_response_format(
         Complete response_format dict with type, json_schema, name, schema, strict
     """
     # Sanitize the schema name
+    schema = pydantic_class.model_json_schema()
     original_name = schema.get("title", pydantic_class.__name__)
     sanitized_name = _sanitize_schema_name(original_name)
 
@@ -138,7 +139,6 @@ def _build_json_schema_response_format(
     # which requires additionalProperties with a schema (for dynamic keys).
     # OpenAI's strict mode only allows additionalProperties: false, so we disable
     # strict mode for these models.
-    from sotopia.database.base_models import LLMEvalBaseModel
 
     use_strict = not issubclass(pydantic_class, LLMEvalBaseModel)
 
@@ -246,10 +246,8 @@ async def agenerate(
             output_parser, PydanticOutputParser
         ), "structured output only supported in PydanticOutputParser"
 
-        # # Build JSON schema response format with OpenAI compatibility fixes
-        schema = output_parser.pydantic_object.model_json_schema()
         response_format = _build_json_schema_response_format(
-            schema, output_parser.pydantic_object
+            output_parser.pydantic_object
         )
 
         # Build completion kwargs with structured output
diff --git a/tests/database/test_evaluation_dimensions.py b/tests/database/test_evaluation_dimensions.py
index c319a58a4..0b4d65cbd 100644
--- a/tests/database/test_evaluation_dimensions.py
+++ b/tests/database/test_evaluation_dimensions.py
@@ -57,13 +57,16 @@ def test_dimension_list() -> Generator[None, None, None]:
 def test_build_dimension_model(test_dimension: Callable[[], None]) -> None:
     # Test building model from dimension id
     model = EvaluationDimensionBuilder.build_dimension_model(["tmppk_test_dimension"])
-    instance = model(test_dimension=("example", 5))
-    assert instance.dict()["test_dimension"] == ("example", 5)
+    # The model expects CustomReasoningScore objects with reasoning and score fields
+    # Pass as dictionary since the class created in the model is different from the one we create here
+    instance = model(test_dimension={"reasoning": "example", "score": 5})
+    assert instance.model_dump()["test_dimension"]["reasoning"] == "example"
+    assert instance.model_dump()["test_dimension"]["score"] == 5
     # Test validation errors for out of range values
     with pytest.raises(ValidationError):
-        model(test_dimension=("example", 11))
+        model(test_dimension={"reasoning": "example", "score": 11})
     with pytest.raises(ValidationError):
-        model(test_dimension=("example", -1))
+        model(test_dimension={"reasoning": "example", "score": -1})
 
 
 def test_build_dimension_model_from_dict() -> None:
@@ -77,12 +80,14 @@ def test_build_dimension_model_from_dict() -> None:
         }
     ]
     model = EvaluationDimensionBuilder.build_dimension_model_from_dict(dimensions)
-
-    instance = model(test_dim=("example", 5))
-    assert instance.dict()["test_dim"] == ("example", 5)
+    # The model expects CustomReasoningScore objects with reasoning and score fields
+    # Pass as dictionary since the class created in the model is different from the one we create here
+    instance = model(test_dim={"reasoning": "example", "score": 5})
+    assert instance.model_dump()["test_dim"]["reasoning"] == "example"
+    assert instance.model_dump()["test_dim"]["score"] == 5
 
     with pytest.raises(ValidationError):
-        model(test_dim=("example", 11))
+        model(test_dim={"reasoning": "example", "score": 11})
 
 
 def test_select_existing_dimension_model_by_name(
@@ -92,11 +97,14 @@ def test_select_existing_dimension_model_by_name(
     model = EvaluationDimensionBuilder.select_existing_dimension_model_by_name(
         ["test_dimension"]
     )
-    instance = model(test_dimension=("example", 5))
-    assert instance.dict()["test_dimension"] == ("example", 5)
+    # The model expects CustomReasoningScore objects with reasoning and score fields
+    # Pass as dictionary since the class created in the model is different from the one we create here
+    instance = model(test_dimension={"reasoning": "example", "score": 5})
+    assert instance.model_dump()["test_dimension"]["reasoning"] == "example"
+    assert instance.model_dump()["test_dimension"]["score"] == 5
 
     with pytest.raises(ValidationError):
-        model(test_dimension=("example", 11))
+        model(test_dimension={"reasoning": "example", "score": 11})
 
 
 def test_select_existing_dimension_model_by_list_name(
@@ -106,8 +114,11 @@ def test_select_existing_dimension_model_by_list_name(
     model = EvaluationDimensionBuilder.select_existing_dimension_model_by_list_name(
         "test_list"
     )
-    instance = model(test_dimension=("example", 5))
-    assert instance.dict()["test_dimension"] == ("example", 5)
+    # The model expects CustomReasoningScore objects with reasoning and score fields
+    # Pass as dictionary since the class created in the model is different from the one we create here
+    instance = model(test_dimension={"reasoning": "example", "score": 5})
+    assert instance.model_dump()["test_dimension"]["reasoning"] == "example"
+    assert instance.model_dump()["test_dimension"]["score"] == 5
 
     with pytest.raises(ValidationError):
-        model(test_dimension=("example", 11))
+        model(test_dimension={"reasoning": "example", "score": 11})
diff --git a/tests/envs/test_evaluators.py b/tests/envs/test_evaluators.py
index d641f085a..a3a774dcb 100644
--- a/tests/envs/test_evaluators.py
+++ b/tests/envs/test_evaluators.py
@@ -21,8 +21,8 @@ def test_rule_based_terminated_evaluator() -> None:
     response = evaluator(
         1,
         [
-            ("Alice", AgentAction(action_type="leave", argument="")),
-            ("Bob", AgentAction(action_type="none", argument="")),
+            ("Alice", AgentAction(action_type="leave", argument="", to=[])),
+            ("Bob", AgentAction(action_type="none", argument="", to=[])),
         ],
     )
     comment = response[0][1][1]
@@ -30,8 +30,8 @@ def test_rule_based_terminated_evaluator() -> None:
     response = evaluator(
         1,
         [
-            ("Alice", AgentAction(action_type="speak", argument="Leave!")),
-            ("Bob", AgentAction(action_type="leave", argument="")),
+            ("Alice", AgentAction(action_type="speak", argument="Leave!", to=[])),
+            ("Bob", AgentAction(action_type="leave", argument="", to=[])),
         ],
     )
     comment = response[0][1][1]
@@ -39,8 +39,8 @@ def test_rule_based_terminated_evaluator() -> None:
     response = evaluator(
         3,
         [
-            ("Alice", AgentAction(action_type="none", argument="")),
-            ("Bob", AgentAction(action_type="none", argument="")),
+            ("Alice", AgentAction(action_type="none", argument="", to=[])),
+            ("Bob", AgentAction(action_type="none", argument="", to=[])),
         ]
         * 3,
     )
@@ -95,8 +95,8 @@ async def test_rule_based_terminated_evaluator_async() -> None:
     response = await evaluator.__acall__(
         1,
         [
-            ("Alice", AgentAction(action_type="leave", argument="")),
-            ("Bob", AgentAction(action_type="none", argument="")),
+            ("Alice", AgentAction(action_type="leave", argument="", to=[])),
+            ("Bob", AgentAction(action_type="none", argument="", to=[])),
         ],
     )
     comment = response[0][1][1]
@@ -104,8 +104,8 @@ async def test_rule_based_terminated_evaluator_async() -> None:
     response = await evaluator.__acall__(
         1,
         [
-            ("Alice", AgentAction(action_type="speak", argument="Leave!")),
-            ("Bob", AgentAction(action_type="leave", argument="")),
+            ("Alice", AgentAction(action_type="speak", argument="Leave!", to=[])),
+            ("Bob", AgentAction(action_type="leave", argument="", to=[])),
         ],
     )
     comment = response[0][1][1]
@@ -113,8 +113,8 @@ async def test_rule_based_terminated_evaluator_async() -> None:
     response = await evaluator.__acall__(
         3,
         [
-            ("Alice", AgentAction(action_type="none", argument="")),
-            ("Bob", AgentAction(action_type="none", argument="")),
+            ("Alice", AgentAction(action_type="none", argument="", to=[])),
+            ("Bob", AgentAction(action_type="none", argument="", to=[])),
         ]
         * 3,
     )
@@ -137,10 +137,10 @@ async def test_reach_goal_llm_evaluator_async(
 ) -> None:
     background = ScriptBackground(
         scenario="Conversation between two friends at a trivia night",
-        agent_names=["Samuel Anderson", "Giselle Rousseau"],
+        agent_names=["Alice", "Bob"],
         agent_backgrounds=[
-            "Samuel Anderson is a 29-year-old male software developer. He/him pronouns. Samuel Anderson can cook very well. Personality and values description: Samuel Anderson, though somewhat impulsive and free-spirited, values enjoyment. His decision-making is often spontaneous, staying within familiar boundaries. Samuel's secrets: He was once a competitive figure skater.",
-            "Giselle Rousseau is a 21-year-old nonbinary art student. They/them pronouns. Giselle Rousseau enjoys biking and photography. Personality and values description: Giselle Rousseau, open-minded and outgoing yet sensitive, advocates care and fairness. Her decision-making is intuitive and inclusive. Giselle's secrets: Sells forged paintings to wealthy clients",
+            "Alice is a 29-year-old female software developer. She/her pronouns. Alice can cook very well. Personality and values description: Alice, though somewhat impulsive and free-spirited, values enjoyment. Her decision-making is often spontaneous, staying within familiar boundaries. Alice's secrets: She was once a competitive figure skater.",
+            "Bob is a 21-year-old male software developer. He/him pronouns. Bob enjoys biking and photography. Personality and values description: Bob, open-minded and outgoing yet sensitive, advocates care and fairness. His decision-making is intuitive and inclusive. Bob's secrets: He was once a competitive figure skater",
         ],
         agent_goals=[
             "Greet your friends and be polite",
@@ -151,16 +151,23 @@ async def test_reach_goal_llm_evaluator_async(
     messages = [
         ("Environment", background),
         ("Environment", SimpleMessage(message="Turn #1")),
-        ("Alice", AgentAction(action_type="speak", argument="Thank you so much!")),
+        (
+            "Alice",
+            AgentAction(action_type="speak", argument="Thank you so much!", to=["Bob"]),
+        ),
         ("Environment", SimpleMessage(message="Turn #2")),
-        ("Bob", AgentAction(action_type="speak", argument="Fuck you!")),
+        ("Bob", AgentAction(action_type="speak", argument="Fuck you!", to=["Alice"])),
         ("Environment", SimpleMessage(message="Turn #3")),
         (
             "Alice",
-            AgentAction(action_type="speak", argument="Hope you have a great weekend."),
+            AgentAction(
+                action_type="speak",
+                argument="Hope you have a great weekend.",
+                to=["Bob"],
+            ),
         ),
         ("Environment", SimpleMessage(message="Turn #4")),
-        ("Bob", AgentAction(action_type="leave", argument="Leave")),
+        ("Bob", AgentAction(action_type="leave", argument="Leave", to=["Alice"])),
     ]
 
     evaluator = EpisodeLLMEvaluator(
diff --git a/tests/envs/test_parallel.py b/tests/envs/test_parallel.py
index d71a7ec1e..79de86bf7 100644
--- a/tests/envs/test_parallel.py
+++ b/tests/envs/test_parallel.py
@@ -13,7 +13,7 @@
 async def test_parallel_sotopia_env() -> None:
     env_profile = EnvironmentProfile(
         pk="test_pk",
-        code_name="test",
+        codename="test",
         scenario="test",
         agent_goals=["test 1", "test 2"],
     )
@@ -51,7 +51,7 @@ async def test_parallel_sotopia_env() -> None:
         max_steps -= 1
         # Create actual AgentAction objects instead of using sample()
         actions = {
-            agent: AgentAction(action_type="speak", argument="test message")
+            agent: AgentAction(action_type="speak", argument="test message", to=[])
             for agent in env.agents
         }
         await env.astep(actions)
@@ -63,7 +63,7 @@ async def test_parallel_sotopia_env() -> None:
 async def test_parallel_sotopia_env_script_writing_single_step() -> None:
     env_profile = EnvironmentProfile(
         pk="test_pk",
-        code_name="test",
+        codename="test",
         scenario="test",
         agent_goals=["test 1", "test 2"],
     )
@@ -104,7 +104,7 @@ async def test_parallel_sotopia_env_script_writing_single_step() -> None:
         max_steps -= 1
         # Create actual AgentAction objects instead of using sample()
         actions = {
-            agent: AgentAction(action_type="speak", argument="test message")
+            agent: AgentAction(action_type="speak", argument="test message", to=[])
             for agent in env.agents
         }
         await env.astep(actions)
@@ -120,7 +120,7 @@ async def test_parallel_sotopia_env_multi_agents_private_messages() -> None:
     """
     env_profile = EnvironmentProfile(
         pk="test_pk",
-        code_name="test",
+        codename="test",
         scenario="test",
         agent_goals=["test 1", "test 2", "test 3"],
     )
@@ -171,7 +171,7 @@ async def test_parallel_sotopia_env_multi_agents_private_messages() -> None:
             argument="psst only for agent2",
             to=["agent2"],
         ),
-        "agent2": AgentAction(action_type="speak", argument="hi all"),
+        "agent2": AgentAction(action_type="speak", argument="hi all", to=[]),
         "agent3": AgentAction(
             action_type="speak",
             argument="psst only for agent1",
diff --git a/tests/integration/test_benchmark.py b/tests/integration/test_benchmark.py
index 880e0fde9..9bd6ee0fa 100644
--- a/tests/integration/test_benchmark.py
+++ b/tests/integration/test_benchmark.py
@@ -306,6 +306,9 @@ def test_sotopia_benchmark(
         )
         mock_initialize_benchmark_combo.assert_called_once_with("")
 
+        # Reset the mock for the second call
+        mock_initialize_benchmark_combo.reset_mock()
+
         benchmark(
             models=[model_name],
             partner_model="not_test_model",
@@ -338,15 +341,34 @@ def test_sotopia_benchmark_display() -> None:
         )
 
         target_believability = (7.0, 3.4462887784189147)
+        test_model_key = f"{model_name} (test) eval_model as the evaluator"
         assert np.allclose(
-            displayed_stats["test_model"]["believability"],
+            displayed_stats[test_model_key]["believability"],
             target_believability,
             atol=0.02,
-        ), f"Got {displayed_stats['test_model']['believability']}, expected {target_believability}"
+        ), f"Got {displayed_stats[test_model_key]['believability']}, expected {target_believability}"
 
         benchmark_file = "/tmp/models_vs_not_test_model.jsonl"
         recovered_data = [json.loads(line) for line in open(benchmark_file, "r")]
-        target_content = '{"model_name": "test_model", "SOC [-10, 0]": 7.0, "SEC [-10, 0]": 7.0, "FIN [-5, 5]": 7.0, "REL [-5, 5]": 7.0, "KNO [0, 10]": 7.0, "GOAL [0, 10]": 7.0, "BEL [0, 10]": 7.0}'
+        # Check that the file was created and has the expected structure
+        assert len(recovered_data) > 0, "No data found in benchmark file"
+        # The model_name in the output is the full key from model_rewards_dict
+        expected_model_name = f"{model_name} (test) eval_model as the evaluator"
         assert (
-            str(recovered_data[0]).replace("'", '"') == target_content
-        ), f"Expected {target_content}, but got {recovered_data[0]}"
+            recovered_data[0]["model_name"] == expected_model_name
+        ), f"Expected model_name to be '{expected_model_name}', got {recovered_data[0]['model_name']}"
+        # Verify all expected dimensions are present with correct values
+        expected_dimensions = {
+            "SOC [-10, 0]": 7.0,
+            "SEC [-10, 0]": 7.0,
+            "FIN [-5, 5]": 7.0,
+            "REL [-5, 5]": 7.0,
+            "KNO [0, 10]": 7.0,
+            "GOAL [0, 10]": 7.0,
+            "BEL [0, 10]": 7.0,
+        }
+        for dim, expected_value in expected_dimensions.items():
+            assert dim in recovered_data[0], f"Missing dimension {dim} in output"
+            assert np.allclose(
+                recovered_data[0][dim], expected_value, atol=0.01
+            ), f"Expected {dim} to be {expected_value}, got {recovered_data[0][dim]}"

From 61f7cbf9e929b7d71827d1548b828ac29b5eb468 Mon Sep 17 00:00:00 2001
From: Xuhui Zhou <zhouxuhui2018@gmail.com>
Date: Thu, 18 Dec 2025 23:37:09 -0500
Subject: [PATCH 19/23] fix mypy

---
 examples/generation_api/custom_model.py      |  3 +--
 sotopia/database/evaluation_dimensions.py    | 13 +++++++------
 sotopia/generation_utils/generate.py         | 13 ++++++++-----
 sotopia/generation_utils/output_parsers.py   |  3 ++-
 tests/database/test_evaluation_dimensions.py | 18 +++++++++---------
 5 files changed, 27 insertions(+), 23 deletions(-)

diff --git a/examples/generation_api/custom_model.py b/examples/generation_api/custom_model.py
index 410a5df46..f76fd1b38 100644
--- a/examples/generation_api/custom_model.py
+++ b/examples/generation_api/custom_model.py
@@ -13,7 +13,6 @@
 from sotopia.generation_utils import (
     ListOfIntOutputParser,
     agenerate,
-    default_temperature,
 )
 import logging
 
@@ -27,7 +26,7 @@ async def generate_n_random_numbers(n: int) -> list[int]:
         template="Generate {n} random integer numbers. {format_instructions}",
         input_values={"n": str(n)},
         output_parser=ListOfIntOutputParser(n),
-        temperature=default_temperature(0.0),
+        temperature=0.0,
     )
 
 
diff --git a/sotopia/database/evaluation_dimensions.py b/sotopia/database/evaluation_dimensions.py
index caa9c6504..b516a313b 100644
--- a/sotopia/database/evaluation_dimensions.py
+++ b/sotopia/database/evaluation_dimensions.py
@@ -7,7 +7,7 @@
     cast,
 )
 
-from pydantic import BaseModel, create_model, field_validator
+from pydantic import create_model, field_validator
 from redis_om import JsonModel
 from redis_om.model.model import Field
 
@@ -267,7 +267,7 @@ def validator(value: tuple[str, int]) -> tuple[str, int]:
     def create_reasoning_score_class(low: int, high: int) -> Type[LLMBaseModel]:
         """Create a custom reasoning score class for the given range"""
 
-        def create_validator(low_val: int, high_val: int):
+        def create_validator(low_val: int, high_val: int) -> Callable[[int], int]:
             def validator(v: int) -> int:
                 if not low_val <= v <= high_val:
                     raise ValueError(f"Score must be between {low_val} and {high_val}")
@@ -284,7 +284,8 @@ class CustomReasoningScore(LLMBaseModel):
             @field_validator("score")
             @classmethod
             def validate_score(cls, v: int) -> int:
-                return create_validator(low, high)(v)
+                validator_func = create_validator(low, high)
+                return validator_func(v)
 
         return CustomReasoningScore
 
@@ -340,7 +341,7 @@ def build_dimension_model_from_dict(
 
         dimension_model = create_model(
             "CustomEvaluationDimensionModel",
-            __base__=BaseModel,
+            __base__=LLMBaseModel,
             **fields,
         )
         return dimension_model
@@ -373,9 +374,9 @@ def select_existing_dimension_model_by_name(
                 Field(..., description=dimension.description),
             )
 
-        model: Type[BaseModel] = create_model(
+        model = create_model(
             "CustomEvaluationDimensionModel",
-            __base__=BaseModel,
+            __base__=LLMBaseModel,
             **fields,
         )
         return model
diff --git a/sotopia/generation_utils/generate.py b/sotopia/generation_utils/generate.py
index 22c44d527..0962a3ee6 100644
--- a/sotopia/generation_utils/generate.py
+++ b/sotopia/generation_utils/generate.py
@@ -5,7 +5,7 @@
 from litellm.litellm_core_utils.get_supported_openai_params import (
     get_supported_openai_params,
 )
-from typing import Any, cast
+from typing import Any
 
 import gin
 
@@ -90,7 +90,10 @@ async def format_bad_output(
             model=model_name,
             messages=[{"role": "user", "content": content}],
         )
-    return response.choices[0].message.content
+    content = response.choices[0].message.content
+    if content is None:
+        raise ValueError("Response content is None")
+    return content
 
 
 def _sanitize_schema_name(name: str) -> str:
@@ -225,7 +228,7 @@ async def agenerate(
         base_url = None
         api_key = None
 
-    temperature_value = cast(float | None, temperature)
+    temperature_value: float | None = temperature
 
     supported_params: list[str] | None = None
     if base_url is None:
@@ -484,7 +487,7 @@ async def agenerate_script(
                     history=history,
                     agent=agent_name,
                 ),
-                output_parser=ScriptOutputParser(  # type: ignore[arg-type]
+                output_parser=ScriptOutputParser(  # type: ignore[call-arg,arg-type]
                     agent_names=agent_names,
                     background=background.to_natural_language(),
                     single_turn=True,
@@ -508,7 +511,7 @@ async def agenerate_script(
                 input_values=dict(
                     background=background.to_natural_language(),
                 ),
-                output_parser=ScriptOutputParser(  # type: ignore[arg-type]
+                output_parser=ScriptOutputParser(  # type: ignore[call-arg,arg-type]
                     agent_names=agent_names,
                     background=background.to_natural_language(),
                     single_turn=False,
diff --git a/sotopia/generation_utils/output_parsers.py b/sotopia/generation_utils/output_parsers.py
index 4a4d0cf4a..4a33bb52e 100644
--- a/sotopia/generation_utils/output_parsers.py
+++ b/sotopia/generation_utils/output_parsers.py
@@ -61,7 +61,8 @@ def extract_value(obj: dict[str, Any] | list[Any] | str) -> Any:
             return self.pydantic_object.model_validate(data, context=context)
         else:
             # Fallback to JSON validation for backward compatibility
-            if "properties" in json_result:
+            # Type narrowing: check that json_result is a dict before accessing "properties"
+            if isinstance(json_result, dict) and "properties" in json_result:
                 return self.pydantic_object.model_validate_json(
                     json.dumps(json_result["properties"])
                 )
diff --git a/tests/database/test_evaluation_dimensions.py b/tests/database/test_evaluation_dimensions.py
index 0b4d65cbd..09a2e2d11 100644
--- a/tests/database/test_evaluation_dimensions.py
+++ b/tests/database/test_evaluation_dimensions.py
@@ -59,14 +59,14 @@ def test_build_dimension_model(test_dimension: Callable[[], None]) -> None:
     model = EvaluationDimensionBuilder.build_dimension_model(["tmppk_test_dimension"])
     # The model expects CustomReasoningScore objects with reasoning and score fields
     # Pass as dictionary since the class created in the model is different from the one we create here
-    instance = model(test_dimension={"reasoning": "example", "score": 5})
+    instance = model(test_dimension={"reasoning": "example", "score": 5})  # type: ignore[call-arg]
     assert instance.model_dump()["test_dimension"]["reasoning"] == "example"
     assert instance.model_dump()["test_dimension"]["score"] == 5
     # Test validation errors for out of range values
     with pytest.raises(ValidationError):
-        model(test_dimension={"reasoning": "example", "score": 11})
+        model(test_dimension={"reasoning": "example", "score": 11})  # type: ignore[call-arg]
     with pytest.raises(ValidationError):
-        model(test_dimension={"reasoning": "example", "score": -1})
+        model(test_dimension={"reasoning": "example", "score": -1})  # type: ignore[call-arg]
 
 
 def test_build_dimension_model_from_dict() -> None:
@@ -82,12 +82,12 @@ def test_build_dimension_model_from_dict() -> None:
     model = EvaluationDimensionBuilder.build_dimension_model_from_dict(dimensions)
     # The model expects CustomReasoningScore objects with reasoning and score fields
     # Pass as dictionary since the class created in the model is different from the one we create here
-    instance = model(test_dim={"reasoning": "example", "score": 5})
+    instance = model(test_dim={"reasoning": "example", "score": 5})  # type: ignore[call-arg]
     assert instance.model_dump()["test_dim"]["reasoning"] == "example"
     assert instance.model_dump()["test_dim"]["score"] == 5
 
     with pytest.raises(ValidationError):
-        model(test_dim={"reasoning": "example", "score": 11})
+        model(test_dim={"reasoning": "example", "score": 11})  # type: ignore[call-arg]
 
 
 def test_select_existing_dimension_model_by_name(
@@ -99,12 +99,12 @@ def test_select_existing_dimension_model_by_name(
     )
     # The model expects CustomReasoningScore objects with reasoning and score fields
     # Pass as dictionary since the class created in the model is different from the one we create here
-    instance = model(test_dimension={"reasoning": "example", "score": 5})
+    instance = model(test_dimension={"reasoning": "example", "score": 5})  # type: ignore[call-arg]
     assert instance.model_dump()["test_dimension"]["reasoning"] == "example"
     assert instance.model_dump()["test_dimension"]["score"] == 5
 
     with pytest.raises(ValidationError):
-        model(test_dimension={"reasoning": "example", "score": 11})
+        model(test_dimension={"reasoning": "example", "score": 11})  # type: ignore[call-arg]
 
 
 def test_select_existing_dimension_model_by_list_name(
@@ -116,9 +116,9 @@ def test_select_existing_dimension_model_by_list_name(
     )
     # The model expects CustomReasoningScore objects with reasoning and score fields
     # Pass as dictionary since the class created in the model is different from the one we create here
-    instance = model(test_dimension={"reasoning": "example", "score": 5})
+    instance = model(test_dimension={"reasoning": "example", "score": 5})  # type: ignore[call-arg]
     assert instance.model_dump()["test_dimension"]["reasoning"] == "example"
     assert instance.model_dump()["test_dimension"]["score"] == 5
 
     with pytest.raises(ValidationError):
-        model(test_dimension={"reasoning": "example", "score": 11})
+        model(test_dimension={"reasoning": "example", "score": 11})  # type: ignore[call-arg]

From 3255c974549385e8972497aa78235c83b6c6ab30 Mon Sep 17 00:00:00 2001
From: Xuhui Zhou <zhouxuhui2018@gmail.com>
Date: Fri, 19 Dec 2025 15:22:46 -0500
Subject: [PATCH 20/23] fix benchmark tests

---
 pyproject.toml                     |   5 +-
 sotopia/cli/benchmark/benchmark.py | 133 ++++++++++++++++++++---------
 2 files changed, 96 insertions(+), 42 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 1132818aa..14e501206 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -48,8 +48,8 @@ realtime = [
     "pyaudio>=0.2.14,<0.3.0",
 ]
 
-[tool.uv]
-dev-dependencies = [
+[dependency-groups]
+dev = [
     "pre-commit",
     "nbmake",
     "types-setuptools",
@@ -85,6 +85,7 @@ aact = { git = "https://github.com/ProKil/aact" , branch = "main" }
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 python_files = "test_*.py"
+asyncio_default_fixture_loop_scope = "function"
 markers = [
     "real_llm: marks tests as requiring real LLM API calls (deselect with '-m \"not real_llm\"')",
     "slow: marks tests as slow running (deselect with '-m \"not slow\"')",
diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py
index 57784c413..0a8b7ef4d 100644
--- a/sotopia/cli/benchmark/benchmark.py
+++ b/sotopia/cli/benchmark/benchmark.py
@@ -519,47 +519,22 @@ def save_to_jsonl(
     print(f"Output saved to {output_fn}")
 
 
-@app.command()
-def benchmark(
-    models: Annotated[
-        list[str],
-        typer.Argument(
-            help=f"All the language model you want to benchmark. Default is the pre-loaded model list {default_model_list}."
-        ),
-    ] = default_model_list,
-    agent_class: Annotated[
-        type[LLMAgent],
-        typer.Argument(
-            help="The agent class you want to use for the evaluated models. Must be a subclass of LLMAgent."
-        ),
-    ] = LLMAgent,
-    partner_model: Annotated[
-        str, typer.Argument(help="The partner model you want to use.")
-    ] = "together_ai/meta-llama/Llama-3-70b-chat-hf",
-    evaluator_model: Annotated[
-        str, typer.Argument(help="The evaluator model you want to use.")
-    ] = "gpt-4o",
-    batch_size: Annotated[
-        int, typer.Argument(help="The batch size you want to use.")
-    ] = 10,
-    task: Annotated[
-        str, typer.Argument(help="The task id you want to benchmark.")
-    ] = "hard",
-    url: Annotated[
-        str, typer.Argument(help="The url to fetch the benchmark combo.")
-    ] = "",
-    print_logs: Annotated[bool, typer.Argument(help="Print logs.")] = False,
-    only_show_performance: Annotated[
-        bool, typer.Argument(help="Only show performance.")
-    ] = False,
-    output_to_jsonl: Annotated[bool, typer.Argument(help="Output to jsonl.")] = False,
-    push_to_db: Annotated[bool, typer.Argument(help="Push to db.")] = False,
-    save_dir: Annotated[
-        str, typer.Argument(help="The directory to save the output.")
-    ] = ".",
-    tag: Annotated[str, typer.Argument(help="The tag for the experiment.")] = "",
+def _benchmark_impl(
+    models: list[str] = default_model_list,
+    agent_class: type[LLMAgent] = LLMAgent,
+    partner_model: str = "together_ai/meta-llama/Llama-3-70b-chat-hf",
+    evaluator_model: str = "gpt-4o",
+    batch_size: int = 10,
+    task: str = "hard",
+    url: str = "",
+    print_logs: bool = False,
+    only_show_performance: bool = False,
+    output_to_jsonl: bool = False,
+    push_to_db: bool = False,
+    save_dir: str = ".",
+    tag: str = "",
 ) -> None:
-    """A simple command-line interface example."""
+    """Internal implementation of benchmark logic."""
     _set_up_logs(print_logs=print_logs)
     benchmark_combo = initialize_benchmark_combo(url)
     if task == "hard":
@@ -682,3 +657,81 @@ def benchmark(
             agent_class=agent_class.__name__,
             tag=tag,
         )
+
+
+@app.command()
+def benchmark(
+    models: Annotated[
+        list[str] | None,
+        typer.Option(
+            "--models",
+            "-m",
+            help=f"Language models to benchmark (default: {default_model_list})",
+        ),
+    ] = None,
+    partner_model: Annotated[
+        str,
+        typer.Option(help="The partner model you want to use."),
+    ] = "together_ai/meta-llama/Llama-3-70b-chat-hf",
+    evaluator_model: Annotated[
+        str,
+        typer.Option(help="The evaluator model you want to use."),
+    ] = "gpt-4o",
+    batch_size: Annotated[
+        int,
+        typer.Option(help="The batch size you want to use."),
+    ] = 10,
+    task: Annotated[
+        str,
+        typer.Option(help="The task id you want to benchmark."),
+    ] = "hard",
+    url: Annotated[
+        str,
+        typer.Option(help="The url to fetch the benchmark combo."),
+    ] = "",
+    print_logs: Annotated[
+        bool,
+        typer.Option(help="Print logs."),
+    ] = False,
+    only_show_performance: Annotated[
+        bool,
+        typer.Option(help="Only show performance."),
+    ] = False,
+    output_to_jsonl: Annotated[
+        bool,
+        typer.Option(help="Output to jsonl."),
+    ] = False,
+    push_to_db: Annotated[
+        bool,
+        typer.Option(help="Push to db."),
+    ] = False,
+    save_dir: Annotated[
+        str,
+        typer.Option(help="The directory to save the output."),
+    ] = ".",
+    tag: Annotated[
+        str,
+        typer.Option(help="The tag for the experiment."),
+    ] = "",
+) -> None:
+    """Run sotopia benchmark using LLMAgent (CLI interface)."""
+    # Handle default for models (can't use default with list[str] in typer.Option)
+    if models is None:
+        models = default_model_list
+
+    # Call the implementation with LLMAgent hard-coded
+    _benchmark_impl(
+        models=models,
+        agent_class=LLMAgent,  # Hard-coded for CLI
+        partner_model=partner_model,
+        evaluator_model=evaluator_model,
+        batch_size=batch_size,
+        task=task,
+        url=url,
+        print_logs=print_logs,
+        only_show_performance=only_show_performance,
+        output_to_jsonl=output_to_jsonl,
+        push_to_db=push_to_db,
+        save_dir=save_dir,
+        tag=tag,
+    )

From 6612c92011f9679909b0d491c727c88e44a731c9 Mon Sep 17 00:00:00 2001
From: Xuhui Zhou <zhouxuhui2018@gmail.com>
Date: Fri, 19 Dec 2025 16:08:57 -0500
Subject: [PATCH 21/23] fix benchmark

---
 tests/integration/test_benchmark.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/integration/test_benchmark.py b/tests/integration/test_benchmark.py
index 9bd6ee0fa..d4ffbae92 100644
--- a/tests/integration/test_benchmark.py
+++ b/tests/integration/test_benchmark.py
@@ -6,7 +6,7 @@
 from unittest.mock import patch
 
 from sotopia.cli.benchmark.benchmark import (
-    benchmark,
+    _benchmark_impl,
     benchmark_display,
     run_async_benchmark_in_batch,
 )
@@ -295,7 +295,7 @@ def test_sotopia_benchmark(
         ), f"Expected 20 episodes in the database, but got {len(EpisodeLog.find().all())}"
 
         # `output_to_jsonl` will be tested in the next test, `push_to_db` has been tested elsewhere, so only test `only_show_performance`
-        benchmark(
+        _benchmark_impl(
             models=[model_name],
             partner_model="not_test_model",
             evaluator_model="eval_model",
@@ -309,7 +309,7 @@ def test_sotopia_benchmark(
         # Reset the mock for the second call
         mock_initialize_benchmark_combo.reset_mock()
 
-        benchmark(
+        _benchmark_impl(
             models=[model_name],
             partner_model="not_test_model",
             evaluator_model="eval_model",

From deb111dbea4fdef4f0b1b7f9ee497f7c76ea90e2 Mon Sep 17 00:00:00 2001
From: Xuhui Zhou <zhouxuhui2018@gmail.com>
Date: Fri, 19 Dec 2025 16:28:09 -0500
Subject: [PATCH 22/23] empty task

---
 tests/integration/test_benchmark.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/integration/test_benchmark.py b/tests/integration/test_benchmark.py
index d4ffbae92..9efd55f0e 100644
--- a/tests/integration/test_benchmark.py
+++ b/tests/integration/test_benchmark.py
@@ -300,6 +300,7 @@ def test_sotopia_benchmark(
             partner_model="not_test_model",
             evaluator_model="eval_model",
             url="",
+            task="",  # Use empty task to avoid loading EnvironmentList
             only_show_performance=False,
             output_to_jsonl=False,
             push_to_db=False,
@@ -314,6 +315,7 @@ def test_sotopia_benchmark(
             partner_model="not_test_model",
             evaluator_model="eval_model",
             url="",
+            task="",  # Use empty task to avoid loading EnvironmentList
             only_show_performance=True,
             output_to_jsonl=False,
             push_to_db=False,

From b8100ec1d2073f7b11f688484a93a1d2df612b47 Mon Sep 17 00:00:00 2001
From: Xuhui Zhou <zhouxuhui2018@gmail.com>
Date: Fri, 19 Dec 2025 16:43:06 -0500
Subject: [PATCH 23/23] custom agent

---
 docs/pages/examples/benchmark.md | 80 ++++++++++++++++++++++++++++++++
 1 file changed, 80 insertions(+)

diff --git a/docs/pages/examples/benchmark.md b/docs/pages/examples/benchmark.md
index ab489210e..bf1bc4cbf 100644
--- a/docs/pages/examples/benchmark.md
+++ b/docs/pages/examples/benchmark.md
@@ -13,6 +13,86 @@ Currently this script would run over 100 simulations on the Sotopia Hard tasks.
 
 An example script is provided in `scripts/display_benchmark_results.sh`
 
+## Using Customized Agents with Benchmark
+
+The default `sotopia benchmark` command uses `LLMAgent` for all agents. If you want to use a customized agent class (e.g., a subclass of `LLMAgent` with custom behavior), you can create your own benchmark script that calls `_benchmark_impl` directly.
+
+Here's an example of how to create a custom benchmark command that uses a customized agent:
+
+```python
+import typer
+from sotopia.cli.benchmark.benchmark import _benchmark_impl
+from sotopia.agents import LLMAgent
+from typing import Any, Type
+from typing_extensions import Annotated
+
+app = typer.Typer(pretty_exceptions_enable=False)
+
+
+# Define your custom agent class
+class CustomSocialWorldModelAgent(LLMAgent):
+    """Example custom agent that extends LLMAgent with additional functionality."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        # Set default social_world_model_name if not provided
+        if "social_world_model_name" not in kwargs:
+            kwargs["social_world_model_name"] = "gpt-4.1-2025-04-14"
+        super().__init__(*args, **kwargs)
+
+
+@app.command(name="run-custom-benchmark")
+def run_custom_benchmark(
+    models: Annotated[
+        str, typer.Option(help="Comma-separated list of models to benchmark")
+    ] = "gpt-4.1-2025-04-14",
+    partner_model: Annotated[
+        str, typer.Option(help="Partner model to use")
+    ] = "gpt-4o-2024-08-06",
+    experiment_tag: Annotated[
+        str, typer.Option(help="Tag for the benchmark run")
+    ] = "custom_agent_trial",
+    batch_size: Annotated[int, typer.Option(help="Batch size for processing")] = 100,
+    push_to_db: Annotated[
+        bool, typer.Option(help="Whether to push results to database")
+    ] = True,
+    evaluator_model: Annotated[
+        str, typer.Option(help="Model to use for evaluation")
+    ] = "gpt-4o",
+    task: Annotated[str, typer.Option(help="Task difficulty level")] = "hard",
+) -> None:
+    """Run benchmark with custom agent class."""
+    # Call _benchmark_impl with your custom agent class
+    _benchmark_impl(
+        models=models.split(","),
+        agent_class=CustomSocialWorldModelAgent,  # Use your custom agent
+        partner_model=partner_model,
+        evaluator_model=evaluator_model,
+        batch_size=batch_size,
+        task=task,
+        push_to_db=push_to_db,
+        tag=experiment_tag,
+    )
+
+
+if __name__ == "__main__":
+    app()
+```
+
+### Key Points:
+
+1. **Custom Agent Class**: Your custom agent must be a subclass of `LLMAgent` (or another agent class that implements the same interface).
+
+2. **Using `_benchmark_impl`**: The `_benchmark_impl` function accepts an `agent_class` parameter that allows you to specify which agent class to use for the benchmark.
+
+3. **Agent Initialization**: When creating your custom agent, make sure it accepts the same initialization parameters as `LLMAgent` (e.g., `agent_profile`, `model_name`, etc.) and passes them to the parent class.
+
+4. **Running the Custom Benchmark**: Save your script and run it like any other Python script:
+   ```bash
+   python your_custom_benchmark.py run-custom-benchmark --models gpt-4.1-2025-04-14
+   ```
+
+For more information on creating custom agents, see the [Creating your own agents](/concepts/agents#creating-your-own-agents) section.
+
 # Benchmark your model as a evaluator
 
 ```