a64378956a
Some checks failed
Pre-commit / run (ubuntu-latest) (push) Has been cancelled
Deploy Sphinx documentation to Pages / build_en (ubuntu-latest, 3.10) (push) Has been cancelled
Deploy Sphinx documentation to Pages / build_zh (ubuntu-latest, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (macos-15, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (macos-15, 3.11) (push) Has been cancelled
Python Unittest Coverage / test (macos-15, 3.12) (push) Has been cancelled
Python Unittest Coverage / test (ubuntu-latest, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (ubuntu-latest, 3.11) (push) Has been cancelled
Python Unittest Coverage / test (ubuntu-latest, 3.12) (push) Has been cancelled
Python Unittest Coverage / test (windows-latest, 3.10) (push) Has been cancelled
Python Unittest Coverage / test (windows-latest, 3.11) (push) Has been cancelled
Python Unittest Coverage / test (windows-latest, 3.12) (push) Has been cancelled
125 lines
3.8 KiB
Python
125 lines
3.8 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""The structured output models used in the werewolf game."""
|
|
from typing import Literal
|
|
|
|
from pydantic import BaseModel, Field, model_validator
|
|
from agentscope.agent import AgentBase
|
|
|
|
|
|
class DiscussionModel(BaseModel):
|
|
"""The output format for discussion."""
|
|
|
|
reach_agreement: bool = Field(
|
|
description="Whether you have reached an agreement or not",
|
|
)
|
|
|
|
|
|
def get_vote_model(agents: list[AgentBase]) -> type[BaseModel]:
|
|
"""Get the vote model by player names."""
|
|
|
|
class VoteModel(BaseModel):
|
|
"""The vote output format."""
|
|
|
|
vote: Literal[tuple(_.name for _ in agents)] = Field( # type: ignore
|
|
description="The name of the player you want to vote for",
|
|
)
|
|
|
|
return VoteModel
|
|
|
|
|
|
class WitchResurrectModel(BaseModel):
|
|
"""The output format for witch resurrect action."""
|
|
|
|
resurrect: bool = Field(
|
|
description="Whether you want to resurrect the player",
|
|
)
|
|
|
|
|
|
def get_poison_model(agents: list[AgentBase]) -> type[BaseModel]:
|
|
"""Get the poison model by player names."""
|
|
|
|
class WitchPoisonModel(BaseModel):
|
|
"""The output format for witch poison action."""
|
|
|
|
poison: bool = Field(
|
|
description="Do you want to use the poison potion",
|
|
)
|
|
name: Literal[ # type: ignore
|
|
tuple(_.name for _ in agents)
|
|
] | None = Field(
|
|
description="The name of the player you want to poison, if you "
|
|
"don't want to poison anyone, just leave it empty",
|
|
default=None,
|
|
)
|
|
|
|
# pylint: disable=no-self-argument
|
|
@model_validator(mode="before")
|
|
def clear_name_if_no_poison(cls, values: dict) -> dict:
|
|
"""Clear name if no poison is used.
|
|
This is to avoid validation errors when `poison` is false but a
|
|
`name` is provided.
|
|
|
|
Args:
|
|
values (`dict`):
|
|
The input data for the model.
|
|
Returns:
|
|
`dict`:
|
|
The validated and possibly modified data.
|
|
"""
|
|
if isinstance(values, dict) and not values.get("poison"):
|
|
values["name"] = None
|
|
return values
|
|
|
|
return WitchPoisonModel
|
|
|
|
|
|
def get_seer_model(agents: list[AgentBase]) -> type[BaseModel]:
|
|
"""Get the seer model by player names."""
|
|
|
|
class SeerModel(BaseModel):
|
|
"""The output format for seer action."""
|
|
|
|
name: Literal[tuple(_.name for _ in agents)] = Field( # type: ignore
|
|
description="The name of the player you want to check",
|
|
)
|
|
|
|
return SeerModel
|
|
|
|
|
|
def get_hunter_model(agents: list[AgentBase]) -> type[BaseModel]:
|
|
"""Get the hunter model by player agents."""
|
|
|
|
class HunterModel(BaseModel):
|
|
"""The output format for hunter action."""
|
|
|
|
shoot: bool = Field(
|
|
description="Whether you want to use the shooting ability or not",
|
|
)
|
|
name: Literal[ # type: ignore
|
|
tuple(_.name for _ in agents)
|
|
] | None = Field(
|
|
description="The name of the player you want to shoot, if you "
|
|
"don't want to the ability, just leave it empty",
|
|
default=None,
|
|
)
|
|
|
|
# pylint: disable=no-self-argument
|
|
@model_validator(mode="before")
|
|
def clear_name_if_no_shoot(cls, values: dict) -> dict:
|
|
"""If shoot is false, set name to None to skip validation.
|
|
This is to avoid validation errors when `shoot` is false but a
|
|
`name` is provided.
|
|
|
|
Args:
|
|
values (`dict`):
|
|
The input data for the model.
|
|
Returns:
|
|
`dict`:
|
|
The validated and possibly modified data.
|
|
"""
|
|
if isinstance(values, dict) and not values.get("shoot"):
|
|
values["name"] = None
|
|
return values
|
|
|
|
return HunterModel
|