diff --git a/.gitignore b/.gitignore index f6ea22a78..ae5bacc2e 100644 --- a/.gitignore +++ b/.gitignore @@ -166,4 +166,6 @@ redis-data/* sotopia/cli/install/redis-data/* redis-stack-server-*/ + +*.rdb examples/experimental/negotiation_arena/redis-data/* 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 ``` diff --git a/examples/experimental/sotopia_original_replica/llm_agent_sotopia.py b/examples/experimental/sotopia_original_replica/llm_agent_sotopia.py index c856e7404..acf3071df 100644 --- a/examples/experimental/sotopia_original_replica/llm_agent_sotopia.py +++ b/examples/experimental/sotopia_original_replica/llm_agent_sotopia.py @@ -10,7 +10,7 @@ from sotopia.database.persistent_profile import AgentProfile from typing import Any -from sotopia.generation_utils import agenerate, StrOutputParser, custom_temperature +from sotopia.generation_utils import agenerate, StrOutputParser # Check Python version if sys.version_info >= (3, 11): @@ -131,7 +131,7 @@ async def aact(self, obs: Observation) -> AgentAction: "agent_name": self.name, }, output_parser=StrOutputParser(), - temperature=custom_temperature(0.7), + temperature=0.7, ) return AgentAction( 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/examples/minimalist_demo.py b/examples/minimalist_demo.py index f9471d3a3..d43af9e9d 100644 --- a/examples/minimalist_demo.py +++ b/examples/minimalist_demo.py @@ -27,7 +27,7 @@ asyncio.run( run_async_server( model_dict={ - "env": "gpt-4", + "env": "gpt-4o-mini", "agent1": "gpt-4o-mini", "agent2": "gpt-4o-mini", }, 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/agents/llm_agent.py b/sotopia/agents/llm_agent.py index b0dcc8c5e..2d2b491ae 100644 --- a/sotopia/agents/llm_agent.py +++ b/sotopia/agents/llm_agent.py @@ -68,7 +68,7 @@ async def aact(self, obs: Observation) -> AgentAction: ) if len(obs.available_actions) == 1 and "none" in obs.available_actions: - return AgentAction(action_type="none", argument="") + return AgentAction(action_type="none", argument="", to=[]) else: # Use agent names from script_background if available agent_names = ( @@ -84,6 +84,7 @@ async def aact(self, obs: Observation) -> AgentAction: agent=self.agent_name, goal=self.goal, script_like=self.script_like, + structured_output=True, agent_names=agent_names, sender=self.agent_name, ) @@ -167,7 +168,7 @@ def act(self, obs: Observation) -> AgentAction: action_type = obs.available_actions[int(input("Action type: "))] argument = input("Argument: ") - return AgentAction(action_type=action_type, argument=argument) + return AgentAction(action_type=action_type, argument=argument, to=[]) async def aact(self, obs: Observation) -> AgentAction: self.recv_message("Environment", obs) @@ -197,7 +198,7 @@ async def aact(self, obs: Observation) -> AgentAction: else: argument = "" - return AgentAction(action_type=action_type, argument=argument) + return AgentAction(action_type=action_type, argument=argument, to=[]) class Agents(dict[str, BaseAgent[Observation, AgentAction]]): diff --git a/sotopia/agents/redis_agent.py b/sotopia/agents/redis_agent.py index 1be56f128..72e672626 100644 --- a/sotopia/agents/redis_agent.py +++ b/sotopia/agents/redis_agent.py @@ -76,7 +76,7 @@ async def aact( ) ) last_timestamp = sorted_message_list[-1][0] - return AgentAction(action_type="none", argument="") + return AgentAction(action_type="none", argument="", to=[]) else: async with aiohttp.ClientSession() as session: # 1. post observation to the message list @@ -137,7 +137,7 @@ async def aact( f"{self._URL}/lock/{self.session_id}/{self.sender_id}/no%20action", ) self.reset("Someone has left or the conversation is too long.") - return AgentAction(action_type="leave", argument="") + return AgentAction(action_type="leave", argument="", to=[]) action_string = sorted_message_list[-1][2] try: action = AgentAction.model_validate_json(action_string) @@ -149,7 +149,7 @@ async def aact( ) ) return AgentAction( - action_type="speak", argument=sorted_message_list[-1][2] + action_type="speak", argument=sorted_message_list[-1][2], to=[] ) def reset( diff --git a/sotopia/cli/benchmark/benchmark.py b/sotopia/cli/benchmark/benchmark.py index 84195ef2e..0a8b7ef4d 100644 --- a/sotopia/cli/benchmark/benchmark.py +++ b/sotopia/cli/benchmark/benchmark.py @@ -41,6 +41,7 @@ from sotopia.server import run_async_server from ..app import app +from typing import Annotated # date and message only FORMAT = "%(asctime)s - %(levelname)s - %(name)s - %(message)s" @@ -71,10 +72,12 @@ def get_avg_reward( - episodes: list[EpisodeLog], model_name: str + episodes: list[EpisodeLog], + model_name: str, + agent_class: str = "", ) -> 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) """ @@ -84,9 +87,13 @@ def get_avg_reward( avg_reward_dict = {} avg_margin_dict = {} avg_variance_dict = {} + 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[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) else: @@ -206,16 +213,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 +235,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: + 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) @@ -298,6 +318,8 @@ def benchmark_display( task: str = "hard", 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 @@ -307,14 +329,24 @@ 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}") 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, # type: ignore + partner_model, + LLMAgent.__name__, + ) + 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) @@ -330,6 +362,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" @@ -349,6 +382,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" @@ -358,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), ], @@ -369,14 +407,17 @@ 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)) @@ -386,23 +427,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( @@ -430,40 +519,22 @@ def save_to_jsonl( print(f"Output saved to {output_fn}") -@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}.", - ), - 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."), +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: - if only_show_performance: - benchmark_display( - models, - partner_model, - evaluator_model, - task, - output_to_jsonl, - save_dir, - ) - return - - """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": @@ -477,7 +548,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") @@ -486,11 +595,26 @@ 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." ) - tag = f"benchmark_{model}_{partner_model}_{evaluator_model}_{task}_trial0" + 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 == "" + else tag + ) typer.echo(typer.style(f"Tag: {tag}", fg=typer.colors.GREEN, bold=True)) model_names = { "env": evaluator_model, @@ -502,6 +626,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: @@ -517,15 +642,96 @@ 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: - return + break - benchmark_display( - models, - partner_model, - evaluator_model, - task, + benchmark_display( + [model], + partner_model, + evaluator_model, + task, + output_to_jsonl=output_to_jsonl, + 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, ) diff --git a/sotopia/cli/install/redis-data/dump.rdb b/sotopia/cli/install/redis-data/dump.rdb deleted file mode 100644 index a69d20f84..000000000 Binary files a/sotopia/cli/install/redis-data/dump.rdb and /dev/null differ diff --git a/sotopia/database/__init__.py b/sotopia/database/__init__.py index 99232487f..3d06c8efa 100644 --- a/sotopia/database/__init__.py +++ b/sotopia/database/__init__.py @@ -8,6 +8,7 @@ from .aggregate_annotations import map_human_annotations_to_episode_logs from .annotators import Annotator +from .base_models import LLMBaseModel, LLMEvalBaseModel from .env_agent_combo_storage import EnvAgentComboStorage from .evaluation_dimensions import ( BaseCustomEvaluationDimension, @@ -61,6 +62,8 @@ from .waiting_room import MatchingInWaitingRoom __all__ = [ + "LLMBaseModel", + "LLMEvalBaseModel", "AgentProfile", "BaseAgentProfile", "EnvironmentProfile", diff --git a/sotopia/database/base_models.py b/sotopia/database/base_models.py index 88ac713ec..c273f8e5d 100644 --- a/sotopia/database/base_models.py +++ b/sotopia/database/base_models.py @@ -16,6 +16,7 @@ else: pass +from pydantic import BaseModel, ConfigDict from redis_om.model.model import NotFoundError from .storage_backend import get_storage_backend, is_local_backend @@ -23,6 +24,30 @@ T = TypeVar("T") +class LLMBaseModel(BaseModel): + """ + Base model for LLM structured output with extra='forbid' configuration. + + This ensures that all fields are required and no extra fields are allowed, + which is compatible with OpenAI's strict mode requirements. + """ + + model_config = ConfigDict(extra="forbid") + + +class LLMEvalBaseModel(LLMBaseModel): + """ + Base model for LLM evaluation dimensions. + + Models inheriting from this class will automatically use non-strict mode + in structured outputs because they typically use dict[str, T] patterns + which are incompatible with strict mode's requirement for + additionalProperties: false. + """ + + pass + + class LocalQueryResult: """A query result object that mimics redis-om's query interface for local storage.""" @@ -218,6 +243,8 @@ def patch_model_for_local_storage(model_class: Type[T]) -> Type[T]: # Re-export for convenience __all__ = [ + "LLMBaseModel", + "LLMEvalBaseModel", "NotFoundError", "LocalQueryResult", "add_local_storage_methods", diff --git a/sotopia/database/evaluation_dimensions.py b/sotopia/database/evaluation_dimensions.py index 4c54ef5af..b516a313b 100644 --- a/sotopia/database/evaluation_dimensions.py +++ b/sotopia/database/evaluation_dimensions.py @@ -1,19 +1,17 @@ from typing import ( TYPE_CHECKING, - Annotated, Any, Callable, - Tuple, Type, Union, cast, ) -from pydantic import AfterValidator, BaseModel, create_model +from pydantic import create_model, field_validator from redis_om import JsonModel from redis_om.model.model import Field -from .base_models import patch_model_for_local_storage +from .base_models import LLMBaseModel, LLMEvalBaseModel, patch_model_for_local_storage from .storage_backend import is_local_backend @@ -35,21 +33,55 @@ def minus_ten_to_zero(v: int) -> int: return v -class SotopiaDimensionsPlus(BaseModel): +class NegativeReasoningScore(LLMBaseModel): + 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(LLMBaseModel): + 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(LLMBaseModel): + 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(LLMEvalBaseModel): """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 " "2. Analyze whether the actions of the agent align with their character traits (e.g., personality, values, and etc.). Start the analysis with tag . " "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. " @@ -57,19 +89,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, " @@ -77,25 +104,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. " @@ -103,21 +122,16 @@ class SotopiaDimensionsPlus(BaseModel): ) -class SotopiaDimensions(BaseModel): +class SotopiaDimensions(LLMEvalBaseModel): """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 " "2. Analyze whether the actions of the agent align with their character traits (e.g., personality, values, and etc.). Start the analysis with tag . " "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. " @@ -125,19 +139,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, " @@ -145,25 +154,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. " @@ -171,12 +172,10 @@ class SotopiaDimensions(BaseModel): ) -class GoalDimension(BaseModel): +class GoalDimension(LLMBaseModel): """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. " @@ -184,7 +183,7 @@ class GoalDimension(BaseModel): ) -class BaseCustomEvaluationDimension(BaseModel): +class BaseCustomEvaluationDimension(LLMBaseModel): pk: str | None = "" name: str = Field(index=True) description: str = Field(index=True) @@ -215,7 +214,7 @@ def __init__(self, **kwargs: Any): super().__init__(**kwargs) -class BaseCustomEvaluationDimensionList(BaseModel): +class BaseCustomEvaluationDimensionList(LLMBaseModel): pk: str | None = "" name: str = Field(index=True) dimension_pks: list[str] = Field(default_factory=list, index=True) @@ -253,18 +252,45 @@ 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 + ) -> 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 build_dimension_model(dimension_ids: list[str]) -> Type[BaseModel]: + 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) -> 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}") + return v + + return validator + + class CustomReasoningScore(LLMBaseModel): + 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: + validator_func = create_validator(low, high) + return validator_func(v) + + return CustomReasoningScore + + @staticmethod + def build_dimension_model(dimension_ids: list[str]) -> Type[LLMBaseModel]: """ Build an evaluation dimension from existing dimension primary keys. The returned model is a pydantic model that can be used to evaluate the conversation. @@ -273,20 +299,20 @@ 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), ) - model: Type[BaseModel] = create_model( + model: Type[LLMBaseModel] = create_model( "CustomEvaluationDimensionModel", - __base__=BaseModel, + __base__=LLMBaseModel, **fields, ) return model @@ -294,7 +320,7 @@ def build_dimension_model(dimension_ids: list[str]) -> Type[BaseModel]: @staticmethod def build_dimension_model_from_dict( dimensions: list[dict[str, Union[str, int]]], - ) -> Type[BaseModel]: + ) -> Type[LLMBaseModel]: """ Build an evaluation dimension from a dictionary that specifies the parameters of the `CustomEvaluationDimension`. The returned model is a pydantic model that can be used to evaluate the conversation. @@ -302,20 +328,20 @@ 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), ) dimension_model = create_model( "CustomEvaluationDimensionModel", - __base__=BaseModel, + __base__=LLMBaseModel, **fields, ) return dimension_model @@ -323,7 +349,7 @@ def build_dimension_model_from_dict( @staticmethod def select_existing_dimension_model_by_name( dimension_names: list[str], - ) -> Type[BaseModel]: + ) -> Type[LLMBaseModel]: """ Build an evaluation dimension from existing dimension names. For example `['believability', 'goal']` The returned model is a pydantic model that can be used to evaluate the conversation. @@ -337,20 +363,20 @@ 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), ) - model: Type[BaseModel] = create_model( + model = create_model( "CustomEvaluationDimensionModel", - __base__=BaseModel, + __base__=LLMBaseModel, **fields, ) return model @@ -358,7 +384,7 @@ def select_existing_dimension_model_by_name( @staticmethod def select_existing_dimension_model_by_list_name( list_name: str, - ) -> Type[BaseModel]: + ) -> Type[LLMBaseModel]: """ Build an evaluation dimension from existing `CustomEvaluationDimensionList` list names. For example, directly use `sotopia` The returned model is a pydantic model that can be used to evaluate the conversation. diff --git a/sotopia/database/logs.py b/sotopia/database/logs.py index ab5f4b06d..183fd9921 100644 --- a/sotopia/database/logs.py +++ b/sotopia/database/logs.py @@ -48,6 +48,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/database/persistent_profile.py b/sotopia/database/persistent_profile.py index 62b6204b1..8b2b5d3c5 100644 --- a/sotopia/database/persistent_profile.py +++ b/sotopia/database/persistent_profile.py @@ -11,6 +11,7 @@ from redis_om import JsonModel from redis_om.model.model import Field +from . import LLMBaseModel from .base_models import patch_model_for_local_storage from .storage_backend import is_local_backend @@ -68,7 +69,7 @@ def __init__(self, **kwargs: Any): super().__init__(**kwargs) -class BaseEnvironmentProfile(BaseModel): +class BaseEnvironmentProfile(LLMBaseModel): pk: str | None = Field(default_factory=lambda: "") codename: str = Field( default="", @@ -129,7 +130,7 @@ def __init__(self, **kwargs: Any): super().__init__(**kwargs) -class BaseRelationshipProfile(BaseModel): +class BaseRelationshipProfile(LLMBaseModel): pk: str | None = Field(default_factory=lambda: "") agent_1_id: str agent_2_id: str diff --git a/sotopia/envs/evaluators.py b/sotopia/envs/evaluators.py index ee471562f..7629db43c 100644 --- a/sotopia/envs/evaluators.py +++ b/sotopia/envs/evaluators.py @@ -11,9 +11,8 @@ from sotopia.generation_utils import ( PydanticOutputParser, agenerate, - custom_temperature, - default_temperature, ) +from sotopia.database import LLMEvalBaseModel from sotopia.messages import ( AgentAction, Message, @@ -25,7 +24,7 @@ T_eval_dim = TypeVar("T_eval_dim", bound=BaseModel) -class EvaluationForAgents(BaseModel, Generic[T_eval_dim]): +class EvaluationForAgents(LLMEvalBaseModel, Generic[T_eval_dim]): evaluations: dict[str, T_eval_dim] @@ -181,12 +180,6 @@ async def __acall__( + "] (no other keys).\n" ) - temperature_setting = ( - default_temperature(temperature) - if temperature == 0.0 - else custom_temperature(temperature) - ) - # Use structured output if model supports it (not just custom/structured endpoints) use_structured_output = self.model_name.startswith( "custom/structured" @@ -204,7 +197,7 @@ async def __acall__( output_parser=PydanticOutputParser[self.response_format_class]( # type: ignore[name-defined] pydantic_object=self.response_format_class ), - temperature=temperature_setting, + temperature=temperature, structured_output=use_structured_output, ) response_list = [] diff --git a/sotopia/envs/parallel.py b/sotopia/envs/parallel.py index 9d24d3b4c..642b306ea 100644 --- a/sotopia/envs/parallel.py +++ b/sotopia/envs/parallel.py @@ -391,7 +391,9 @@ def step( # Masking actions from agent that are in turn for idx, agent in enumerate(self.agents): if not self.action_mask[idx]: - complied_actions[agent] = AgentAction(action_type="none", argument="") + complied_actions[agent] = AgentAction( + action_type="none", argument="", to=[] + ) self.recv_message( "Environment", SimpleMessage(message=f"Turn #{self.turn_number}") @@ -480,7 +482,9 @@ async def astep( # Masking actions from agent that are in turn for idx, agent in enumerate(self.agents): if not self.action_mask[idx]: - complied_actions[agent] = AgentAction(action_type="none", argument="") + complied_actions[agent] = AgentAction( + action_type="none", argument="", to=[] + ) self.recv_message( "Environment", SimpleMessage(message=f"Turn #{self.turn_number}") diff --git a/sotopia/generation_utils/__init__.py b/sotopia/generation_utils/__init__.py index 09fd44a69..857face1e 100644 --- a/sotopia/generation_utils/__init__.py +++ b/sotopia/generation_utils/__init__.py @@ -1,10 +1,7 @@ from .generate import ( - TemperatureSetting, agenerate_env_profile, agenerate, agenerate_action, - custom_temperature, - default_temperature, ) from .output_parsers import ( EnvResponse, @@ -23,7 +20,4 @@ "agenerate_env_profile", "agenerate", "agenerate_action", - "TemperatureSetting", - "default_temperature", - "custom_temperature", ] diff --git a/sotopia/generation_utils/generate.py b/sotopia/generation_utils/generate.py index c6e42af3b..0962a3ee6 100644 --- a/sotopia/generation_utils/generate.py +++ b/sotopia/generation_utils/generate.py @@ -1,14 +1,11 @@ import logging import os -import json -from dataclasses import dataclass from litellm import acompletion -from litellm.exceptions import BadRequestError from litellm.utils import supports_response_schema 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 @@ -16,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, @@ -52,36 +49,19 @@ log.addHandler(console_handler) # subject to future OpenAI changes -DEFAULT_BAD_OUTPUT_PROCESS_MODEL = "gpt-4o-mini" +DEFAULT_BAD_OUTPUT_PROCESS_MODEL = "gpt-5-mini-2025-08-07" DEFAULT_TEMPERATURE = 0.7 -_TEMPERATURE_SENTINEL = object() - -# Cache temperature support per (model, base_url) to avoid repeated bad requests -_TEMPERATURE_SUPPORT_CACHE: dict[tuple[str, str | None], bool] = {} - - -@dataclass(frozen=True) -class TemperatureSetting: - value: float | None - treat_as_default: bool = False - - -def default_temperature(value: float | None) -> TemperatureSetting: - return TemperatureSetting(value=value, treat_as_default=True) - - -def custom_temperature(value: float | None) -> TemperatureSetting: - return TemperatureSetting(value=value, treat_as_default=False) @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, base_url: str | None = None, ) -> 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} @@ -96,43 +76,24 @@ async def format_bad_output( "format_instructions": format_instructions, } content = template.format(**input_values) - - # Build completion kwargs - completion_kwargs: dict[str, Any] = { - "model": model_name, - "messages": [{"role": "user", "content": content}], - } - - # Parse format_instructions to get the schema - try: - schema = json.loads(format_instructions) - # Apply schema fixes for OpenAI compatibility - has_tuples, has_optional_fields = _check_schema_compatibility(schema) - use_strict = not (has_tuples or has_optional_fields) - _apply_schema_fixes(schema) - - # Build proper json_schema response_format - completion_kwargs["response_format"] = { - "type": "json_schema", - "json_schema": { - "name": "reformatted_output", - "schema": schema, - "strict": use_strict, - }, - } - except json.JSONDecodeError: - # Fallback to json_object if schema parsing fails - completion_kwargs["response_format"] = {"type": "json_object"} - - # Add base_url if provided - if base_url is not None: - completion_kwargs["base_url"] = base_url - - response = await acompletion(**completion_kwargs) - reformatted_output = response.choices[0].message.content - assert isinstance(reformatted_output, str) - log.info(f"Reformated output: {reformatted_output}") - return reformatted_output + if isinstance(output_parser, PydanticOutputParser): + response_format = _build_json_schema_response_format( + output_parser.pydantic_object, + ) + response = await acompletion( + model=model_name, + response_format=response_format, + messages=[{"role": "user", "content": content}], + ) + else: + response = await acompletion( + model=model_name, + messages=[{"role": "user", "content": 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: @@ -155,113 +116,8 @@ def _sanitize_schema_name(name: str) -> str: return "".join(c if c.isalnum() or c in ("_", "-") else "_" for c in name) -def _check_schema_compatibility(schema: dict[str, Any]) -> tuple[bool, bool]: - """ - Check if Pydantic schema is compatible with OpenAI strict mode. - - Checks for: - - Tuple types (prefixItems) - not supported in strict mode - - Optional fields (properties not in required) - not supported in strict mode - - Args: - schema: JSON schema from pydantic_object.model_json_schema() - - Returns: - (has_tuples, has_optional_fields) - both must be False for strict mode - """ - has_tuples = False - has_optional = False - - def check_recursive(obj: Any) -> tuple[bool, bool]: - """Recursively check schema structure.""" - nonlocal has_tuples, has_optional - - if isinstance(obj, dict): - # Check for arrays with prefixItems (tuples) - if obj.get("type") == "array" and "prefixItems" in obj: - has_tuples = True - - # Check for optional fields (properties not in required list) - if obj.get("type") == "object" and "properties" in obj: - required = set(obj.get("required", [])) - properties = set(obj["properties"].keys()) - if properties != required: - has_optional = True - - # Recurse into nested structures - for value in obj.values(): - if isinstance(value, dict): - check_recursive(value) - elif isinstance(value, list): - for item in value: - if isinstance(item, dict): - check_recursive(item) - - elif isinstance(obj, list): - for item in obj: - if isinstance(item, dict): - check_recursive(item) - - return has_tuples, has_optional - - return check_recursive(schema) - - -def _apply_schema_fixes(schema: dict[str, Any]) -> None: - """ - Apply in-place fixes to make schema more OpenAI-compatible. - - Fixes applied: - - Add 'items': {} to arrays with prefixItems (OpenAI requirement) - - Add 'additionalProperties': False to all objects (strict mode requirement) - - Args: - schema: JSON schema to modify in-place - - Side effects: - Modifies schema dict in-place - """ - - def fix_recursive(obj: Any) -> None: - """Recursively apply fixes to schema.""" - if isinstance(obj, dict): - # Fix arrays with prefixItems - if ( - obj.get("type") == "array" - and "prefixItems" in obj - and "items" not in obj - ): - # OpenAI requires items field even with prefixItems - # Empty object {} means no additional items allowed - obj["items"] = {} - - # Fix objects - add additionalProperties - if obj.get("type") == "object": - if "additionalProperties" not in obj: - obj["additionalProperties"] = False - elif isinstance(obj.get("additionalProperties"), dict): - # Recurse into dict-like objects - fix_recursive(obj["additionalProperties"]) - - # Recurse into nested structures - for value in obj.values(): - if isinstance(value, dict): - fix_recursive(value) - elif isinstance(value, list): - for item in value: - if isinstance(item, dict): - fix_recursive(item) - - elif isinstance(obj, list): - for item in obj: - if isinstance(item, dict): - fix_recursive(item) - - fix_recursive(schema) - - 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. @@ -277,17 +133,17 @@ 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) - # Check compatibility with strict mode - has_tuples, has_optional_fields = _check_schema_compatibility(schema) + # Determine whether to use strict mode + # LLMEvalBaseModel and its subclasses (like EvaluationForAgents) use dict[str, T] + # 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. - # Use strict mode only if schema doesn't contain tuples or optional fields - use_strict = not (has_tuples or has_optional_fields) - - # Apply schema fixes - _apply_schema_fixes(schema) + use_strict = not issubclass(pydantic_class, LLMEvalBaseModel) # Build response format return { @@ -300,144 +156,6 @@ def _build_json_schema_response_format( } -def _resolve_temperature_setting( - temperature: TemperatureSetting | float | None | object, -) -> tuple[float | None, bool, bool]: - """ - Normalize temperature setting to effective value and flags. - - Args: - temperature: Temperature setting (sentinel, TemperatureSetting, float, or None) - - Returns: - Tuple of (effective_temperature, treat_as_default, user_provided): - - effective_temperature: The actual temperature value to use (None means no temp) - - treat_as_default: True if this should be treated as default (not sent to API) - - user_provided: True if user explicitly provided temperature - """ - if temperature is _TEMPERATURE_SENTINEL: - return DEFAULT_TEMPERATURE, True, False - elif isinstance(temperature, TemperatureSetting): - return temperature.value, temperature.treat_as_default, True - else: - return cast(float | None, temperature), False, True - - -def _should_send_temperature( - cache_key: tuple[str, str | None], - effective_temperature: float | None, - treat_as_default: bool, - user_provided: bool, - supported_params: list[str] | None, - model_name: str, -) -> tuple[bool, bool]: - """ - Determine if temperature should be sent to API and check model support. - - Checks cache and model capabilities to decide if temperature parameter - should be included in the API call. - - Args: - cache_key: (model_name, base_url) tuple for cache lookup - effective_temperature: The temperature value to potentially send - treat_as_default: Whether this is a default value (shouldn't be sent) - user_provided: Whether user explicitly provided temperature - supported_params: List of supported params for the model (None if unknown) - model_name: Model name for logging - - Returns: - Tuple of (send_temperature, non_default_requested): - - send_temperature: Whether to include temperature in API call - - non_default_requested: Whether user requested non-default temp (for warnings) - """ - # Initial check: only send if user provided and not treating as default - send_temperature = ( - user_provided and effective_temperature is not None and not treat_as_default - ) - non_default_requested = send_temperature and not treat_as_default - - if send_temperature: - # Check cache first - if cache_key in _TEMPERATURE_SUPPORT_CACHE: - if not _TEMPERATURE_SUPPORT_CACHE[cache_key]: - send_temperature = False - if non_default_requested: - log.warning( - "Model %s previously rejected temperature; ignoring temperature=%s", - model_name, - effective_temperature, - ) - # Check model capabilities if cache miss and we have param info - elif supported_params is not None and "temperature" not in supported_params: - _TEMPERATURE_SUPPORT_CACHE[cache_key] = False - send_temperature = False - if non_default_requested: - log.warning( - "Model %s does not support temperature; ignoring temperature=%s", - model_name, - effective_temperature, - ) - - return send_temperature, non_default_requested - - -def _build_structured_completion_kwargs( - model_name: str, - messages: list[dict[str, str]], - response_format: dict[str, Any], - base_url: str | None, - api_key: str | None, -) -> dict[str, Any]: - """ - Build completion kwargs for structured output mode. - - Args: - model_name: Model identifier - messages: List of message dicts for the API call - response_format: Response format dict with JSON schema - base_url: Optional base URL for custom models - api_key: Optional API key for custom models - - Returns: - Dict of kwargs ready for litellm acompletion - """ - return dict( - model=model_name, - messages=messages, - response_format=response_format, - drop_params=True, # drop params to avoid model error if the model does not support it - base_url=base_url, - api_key=api_key, - ) - - -def _build_standard_completion_kwargs( - model_name: str, - messages: list[dict[str, str]], - base_url: str | None, - api_key: str | None, -) -> dict[str, Any]: - """ - Build completion kwargs for standard (non-structured) output mode. - - Args: - model_name: Model identifier - messages: List of message dicts for the API call - base_url: Optional base URL for custom models - api_key: Optional API key for custom models - - Returns: - Dict of kwargs ready for litellm acompletion - """ - return dict( - model=model_name, - messages=messages, - drop_params=True, - base_url=base_url, - api_key=api_key, - ) - - @gin.configurable @validate_call async def agenerate( @@ -445,7 +163,7 @@ async def agenerate( template: str, input_values: dict[str, str], output_parser: OutputParser[OutputType], - temperature: TemperatureSetting | float | None | object = _TEMPERATURE_SENTINEL, + temperature: float | None = DEFAULT_TEMPERATURE, structured_output: bool = False, bad_output_process_model: str | None = None, use_fixed_model_version: bool = True, @@ -455,15 +173,16 @@ async def agenerate( Generate text using LiteLLM with optional structured output support. This function handles template formatting, temperature management, schema validation, - and API calls with retry logic. It supports both structured (JSON schema) and - standard text generation modes. + and API calls. It supports both structured (JSON schema) and standard text generation modes. Args: model_name: Model identifier (e.g., "gpt-4o-mini" or "custom/model@http://...") template: Template string with {variable} placeholders input_values: Dict mapping template variables to their values output_parser: Parser to convert raw output to desired type - temperature: Temperature setting (sentinel for default, TemperatureSetting, or float) + temperature: Temperature value (float or None). Defaults to DEFAULT_TEMPERATURE. + None means use provider default. With drop_params=True, litellm + automatically handles unsupported parameters. structured_output: If True, use JSON schema mode (requires PydanticOutputParser) bad_output_process_model: Model to use for reformatting bad outputs (if needed) use_fixed_model_version: Whether to use fixed model versioning @@ -486,6 +205,9 @@ async def agenerate( ... ) """ # 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() @@ -506,51 +228,12 @@ async def agenerate( base_url = None api_key = None - cache_key = (model_name, base_url) + temperature_value: float | None = temperature supported_params: list[str] | None = None if base_url is None: supported_params = get_supported_openai_params(model=model_name) - # Resolve temperature setting - effective_temperature, treat_as_default, user_provided = ( - _resolve_temperature_setting(temperature) - ) - - # Determine if temperature should be sent (checks cache and model support) - send_temperature, non_default_requested = _should_send_temperature( - cache_key=cache_key, - effective_temperature=effective_temperature, - treat_as_default=treat_as_default, - user_provided=user_provided, - supported_params=supported_params, - model_name=model_name, - ) - - async def _call_with_retry(completion_kwargs: dict[str, Any]) -> Any: - nonlocal send_temperature - call_kwargs = dict(completion_kwargs) - if send_temperature: - call_kwargs["temperature"] = effective_temperature - try: - response = await acompletion(**call_kwargs) - if send_temperature: - _TEMPERATURE_SUPPORT_CACHE[cache_key] = True - return response - except BadRequestError as exc: - if send_temperature and "temperature" in str(exc).lower(): - send_temperature = False - _TEMPERATURE_SUPPORT_CACHE[cache_key] = False - if non_default_requested: - log.warning( - "Model %s does not support temperature; ignoring temperature=%s", - model_name, - effective_temperature, - ) - call_kwargs.pop("temperature", None) - return await acompletion(**call_kwargs) - raise - messages = [{"role": "user", "content": template}] if structured_output: if not base_url: @@ -566,28 +249,34 @@ async def _call_with_retry(completion_kwargs: dict[str, Any]) -> Any: 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 ) - completion_kwargs = _build_structured_completion_kwargs( - model_name=model_name, - messages=messages, - response_format=response_format, - base_url=base_url, - api_key=api_key, - ) - response = await _call_with_retry(completion_kwargs) + # Build completion kwargs with structured output + completion_kwargs = { + "model": model_name, + "messages": messages, + "response_format": response_format, + "drop_params": True, # litellm automatically drops unsupported params + "base_url": base_url, + "api_key": api_key, + } + if temperature_value is not None: + completion_kwargs["temperature"] = temperature_value + response = await acompletion(**completion_kwargs) else: - completion_kwargs = _build_standard_completion_kwargs( - model_name=model_name, - messages=messages, - base_url=base_url, - api_key=api_key, - ) - response = await _call_with_retry(completion_kwargs) + # Build completion kwargs for standard output + completion_kwargs = { + "model": model_name, + "messages": messages, + "drop_params": True, # litellm automatically drops unsupported params + "base_url": base_url, + "api_key": api_key, + } + if temperature_value is not None: + completion_kwargs["temperature"] = temperature_value + response = await acompletion(**completion_kwargs) result = response.choices[0].message.content # Only PydanticOutputParser supports context parameter @@ -600,7 +289,7 @@ async def _call_with_retry(completion_kwargs: dict[str, Any]) -> Any: except Exception: 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, base_url=base_url, @@ -620,7 +309,7 @@ async def agenerate_env_profile( model_name: str, inspiration_prompt: str = "asking my boyfriend to stop being friends with his ex", examples: str = "", - temperature: TemperatureSetting | float | None | object = _TEMPERATURE_SENTINEL, + temperature: float | None = DEFAULT_TEMPERATURE, bad_output_process_model: str | None = None, use_fixed_model_version: bool = True, ) -> EnvironmentProfile: @@ -683,10 +372,11 @@ async def agenerate_action( action_types: list[ActionType], agent: str, goal: str, - temperature: TemperatureSetting | float | None | object = _TEMPERATURE_SENTINEL, + temperature: float | None = DEFAULT_TEMPERATURE, script_like: bool = False, bad_output_process_model: str | None = None, use_fixed_model_version: bool = True, + structured_output: bool = False, agent_names: list[str] | None = None, sender: str | None = None, ) -> AgentAction: @@ -748,14 +438,14 @@ async def agenerate_action( ), output_parser=PydanticOutputParser(pydantic_object=AgentAction), temperature=temperature, - structured_output=True, + structured_output=structured_output, bad_output_process_model=bad_output_process_model, use_fixed_model_version=use_fixed_model_version, context=validation_context, ) except Exception as e: log.warning(f"Failed to generate action due to {e}") - return AgentAction(action_type="none", argument="") + return AgentAction(action_type="none", argument="", to=[]) @gin.configurable @@ -763,7 +453,7 @@ async def agenerate_action( async def agenerate_script( model_name: str, background: ScriptBackground, - temperature: TemperatureSetting | float | None | object = _TEMPERATURE_SENTINEL, + temperature: float | None = DEFAULT_TEMPERATURE, agent_names: list[str] = [], agent_name: str = "", history: str = "", @@ -797,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, @@ -821,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 3ac63c101..4a33bb52e 100644 --- a/sotopia/generation_utils/output_parsers.py +++ b/sotopia/generation_utils/output_parsers.py @@ -1,14 +1,16 @@ import json import re -from typing import Any, Generic, Type, TypeVar, Optional +from typing import Generic, Type, TypeVar, Optional, Any from pydantic import BaseModel, Field import json_repair +from sotopia.database import LLMBaseModel + OutputType = TypeVar("OutputType", bound=object) T = TypeVar("T", bound=BaseModel) -class EnvResponse(BaseModel): +class EnvResponse(LLMBaseModel): reasoning: str = Field( description="first reiterate agents' social goals and then reason about what agents say/do and whether that aligns with their goals." ) @@ -16,7 +18,7 @@ class EnvResponse(BaseModel): p2_rate: int = Field(description="rating of participant 2, on the scale of 0 to 9") -class OutputParser(BaseModel, Generic[OutputType]): +class OutputParser(LLMBaseModel, Generic[OutputType]): def parse(self, result: str) -> OutputType: raise NotImplementedError @@ -36,9 +38,21 @@ def parse(self, result: str, context: dict[str, Any] | None = None) -> T: json_result = json_repair.loads(result) assert isinstance(json_result, dict) - # Extract the data dict - if "properties" in json_result: - data = json_result["properties"] + # 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"]) + ) else: data = json_result @@ -47,7 +61,8 @@ def parse(self, result: str, context: dict[str, Any] | None = None) -> T: 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/sotopia/messages/message_classes.py b/sotopia/messages/message_classes.py index b6e9a802a..25c644fed 100644 --- a/sotopia/messages/message_classes.py +++ b/sotopia/messages/message_classes.py @@ -1,15 +1,16 @@ import re from typing import Literal, cast -from pydantic import BaseModel, Field, field_validator, ValidationInfo +from pydantic import Field, field_validator, ValidationInfo from pydantic_core import PydanticCustomError +from sotopia.database import LLMBaseModel from sotopia.utils import format_docstring ActionType = Literal["none", "speak", "non-verbal communication", "action", "leave"] -class Message(BaseModel): +class Message(LLMBaseModel): """ An interface for messages. There is only one required method: to_natural_language @@ -131,9 +132,8 @@ class AgentAction(Message): description="the utterance if choose to speak, the expression or gesture if choose non-verbal communication, or the physical action if choose action" ) # New structured fields for private messages - to: list[str] | None = Field( - default=None, - description="recipient name(s), when specified, the action is only visible to the listed agents", + to: list[str] = Field( + description="recipient name(s), when specified, the action is only visible to the listed agents, empty list means public action", ) def to_natural_language(self) -> str: @@ -325,7 +325,7 @@ def parse( name: str = cast(str, res["name"]) parsed_action = AgentAction( - action_type=cast(ActionType, action), argument=argument + action_type=cast(ActionType, action), argument=argument, to=[] ) if name not in agent_names: print( @@ -340,7 +340,7 @@ def parse( f"The error is: {e}", ) raise e - parsed_action = AgentAction(action_type="none", argument="") + parsed_action = AgentAction(action_type="none", argument="", to=[]) name = agent_names[line_idx % 2] # TODO same question as above inactive_agent_name = ( agent_names[0] if name == agent_names[1] else agent_names[1] @@ -365,7 +365,9 @@ def parse( ( inactive_agent_name, "Environment", - AgentAction(action_type="none", argument="did nothing"), + AgentAction( + action_type="none", argument="did nothing", to=[] + ), ), ], ) @@ -458,6 +460,6 @@ def default_value_for_return_type() -> ScriptInteractionReturnType: ] ] results_2: list[tuple[str, Message]] = [ - ("", AgentAction(action_type="none", argument="")) + ("", AgentAction(action_type="none", argument="", to=[])) ] return (results_1, results_2) diff --git a/sotopia/server.py b/sotopia/server.py index 20788c72c..e54111cc2 100644 --- a/sotopia/server.py +++ b/sotopia/server.py @@ -170,7 +170,9 @@ async def generate_messages() -> ( agent_mask = env.action_mask for idx in range(len(agent_mask)): if agent_mask[idx] == 0: - actions[idx] = AgentAction(action_type="none", argument="") + actions[idx] = AgentAction( + action_type="none", argument="", to=[] + ) else: pass @@ -225,7 +227,6 @@ async def generate_messages() -> ( " ".join(info[agent_name]["comments"] for agent_name in env.agents) ) done = all(terminated.values()) - epilog = EpisodeLog( environment=env.profile.pk, agents=[agent.profile.pk for agent in agent_list], diff --git a/tests/database/test_evaluation_dimensions.py b/tests/database/test_evaluation_dimensions.py index c319a58a4..09a2e2d11 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}) # 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=("example", 11)) + model(test_dimension={"reasoning": "example", "score": 11}) # type: ignore[call-arg] with pytest.raises(ValidationError): - model(test_dimension=("example", -1)) + model(test_dimension={"reasoning": "example", "score": -1}) # type: ignore[call-arg] 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}) # 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=("example", 11)) + model(test_dim={"reasoning": "example", "score": 11}) # type: ignore[call-arg] 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}) # 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=("example", 11)) + model(test_dimension={"reasoning": "example", "score": 11}) # type: ignore[call-arg] 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}) # 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=("example", 11)) + model(test_dimension={"reasoning": "example", "score": 11}) # type: ignore[call-arg] 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..9efd55f0e 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,22 +295,27 @@ 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", url="", + task="", # Use empty task to avoid loading EnvironmentList only_show_performance=False, output_to_jsonl=False, push_to_db=False, ) mock_initialize_benchmark_combo.assert_called_once_with("") - benchmark( + # Reset the mock for the second call + mock_initialize_benchmark_combo.reset_mock() + + _benchmark_impl( models=[model_name], 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, @@ -338,15 +343,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]}" diff --git a/uv.lock b/uv.lock index 434674258..72b2916ab 100644 --- a/uv.lock +++ b/uv.lock @@ -3954,7 +3954,7 @@ requires-dist = [ { name = "rich", specifier = ">=13.6.0,<15.0.0" }, { name = "scipy", marker = "extra == 'examples'" }, { name = "streamlit", marker = "extra == 'api'" }, - { name = "together", specifier = ">=0.2.4,<1.5.0" }, + { name = "together", specifier = ">=0.2.4,<1.6.0" }, { name = "torch", marker = "extra == 'examples'" }, { name = "transformers", marker = "extra == 'examples'" }, { name = "uvicorn", marker = "extra == 'api'" },