85 lines
2.5 KiB
Python
85 lines
2.5 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""The structured output models used in the werewolf game."""
|
|
from typing import Literal
|
|
|
|
from agentscope.agent import AgentBase
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
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[tuple(_.name for _ in agents)] | None = Field( # type: ignore
|
|
description="The name of the player you want to poison, if you "
|
|
"don't want to poison anyone, just leave it empty",
|
|
default=None,
|
|
)
|
|
|
|
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[tuple(_.name for _ in agents)] | None = Field( # type: ignore
|
|
description="The name of the player you want to shoot, if you "
|
|
"don't want to the ability, just leave it empty",
|
|
default=None,
|
|
)
|
|
|
|
return HunterModel
|