werewolf/wolf/wolf_agent.py

from agent_build_sdk.model.roles import ROLE_WOLF
from agent_build_sdk.model.werewolf_model import (
    AgentResp, AgentReq, STATUS_START, STATUS_WOLF_SPEECH,
    STATUS_VOTE_RESULT, STATUS_SKILL, STATUS_SKILL_RESULT, STATUS_NIGHT_INFO,
    STATUS_DAY, STATUS_DISCUSS, STATUS_VOTE, STATUS_RESULT, STATUS_NIGHT,
    STATUS_SHERIFF_ELECTION, STATUS_SHERIFF_SPEECH, STATUS_SHERIFF_PK,
    STATUS_SHERIFF_VOTE, STATUS_SHERIFF_SPEECH_ORDER, STATUS_SHERIFF, STATUS_HUNTER, STATUS_HUNTER_RESULT
)
from agent_build_sdk.utils.logger import logger
from agent_build_sdk.sdk.role_agent import BasicRoleAgent
from agent_build_sdk.sdk.agent import format_prompt
from wolf.prompt import (
    DESC_PROMPT, VOTE_PROMPT, KILL_PROMPT, WOLF_SPEECH_PROMPT, GAME_RULE_PROMPT,
    CLEAN_USER_PROMPT, SHERIFF_ELECTION_PROMPT, SHERIFF_SPEECH_PROMPT,
    SHERIFF_VOTE_PROMPT, SHERIFF_SPEECH_ORDER_PROMPT, SHERIFF_TRANSFER_PROMPT
)


class WolfAgent(BasicRoleAgent):
    """Wolf Role Agent"""

    def __init__(self, model_name):
        super().__init__(ROLE_WOLF, model_name=model_name)
        self.memory.set_variable("teammates", [])  # Store teammates information

    def perceive(self, req=AgentReq):
        if req.status == STATUS_START:
            self.memory.clear()
            self.memory.set_variable("name", req.name)
            self.memory.set_variable("teammates", [])  # Reset teammates information
            self.memory.append_history(GAME_RULE_PROMPT)
            self.memory.append_history(f"Host: Hello, your assigned role is [Wolf], you are {req.name}")
            if req.message:  # If there is teammates information
                teammates = req.message.split(",")
                self.memory.set_variable("teammates", teammates)
                self.memory.append_history(f"Host: Your wolf teammates are: {req.message}")
        elif req.status == STATUS_NIGHT:
            self.memory.append_history("Host: Now entering the night, it's dark, please close your eyes")
        elif req.status == STATUS_WOLF_SPEECH:
            # Communication between wolves
            if req.name:
                self.memory.append_history(f"Wolf {req.name} says: {req.message}")
            else:
                self.memory.append_history("Host: Wolves, please open your eyes, confirm each other's identities, and choose your kill target")
        elif req.status == STATUS_SKILL_RESULT:
            self.memory.append_history(f"Host: Wolves, the target you chose to kill tonight is: {req.name}")
        elif req.status == STATUS_NIGHT_INFO:
            self.memory.append_history(f"Host: It's dawn! Last night's information is: {req.message}")
        elif req.status == STATUS_DISCUSS:  # Discussion phase
            if req.name:
                # Other players' speeches
                # Can use model to filter out players' injection messages, or use a smaller model. In actual use, need to consider locking memory to avoid losing messages during interact
                # clean_user_message_prompt = format_prompt(CLEAN_USER_PROMPT, {"user_message": req.message})
                # req.message = self.llm_caller(clean_user_message_prompt)
                self.memory.append_history(req.name + ': ' + req.message)
            else:
                # Host's speech
                self.memory.append_history('Host: Now entering day {}.'.format(str(req.round)))
                self.memory.append_history('Host: Each player describes their information.')
            self.memory.append_history("---------------------------------------------")
        elif req.status == STATUS_VOTE:  # Voting phase
            self.memory.append_history(f'Day {req.round}. Vote information: {req.name} voted for {req.message}')
        elif req.status == STATUS_VOTE_RESULT:  # Voting phase
            if req.name:
                self.memory.append_history(f'Host: The voting result is: {req.name}.')
            else:
                self.memory.append_history('Host: No one is eliminated.')
        elif req.status == STATUS_SHERIFF_ELECTION:
            self.memory.append_history(f"Host: Players running for sheriff: {req.message}")
        elif req.status == STATUS_SHERIFF_SPEECH:
            self.memory.append_history(f"{req.name} (Sheriff campaign speech): {req.message}")
        elif req.status == STATUS_SHERIFF_VOTE:
            self.memory.append_history(f"Sheriff vote: {req.name} voted for {req.message}")
        elif req.status == STATUS_SHERIFF:
            if req.name:
                self.memory.append_history(f"Host: Sheriff badge goes to: {req.name}")
                self.memory.set_variable("sheriff", req.name)
            if req.message:
                self.memory.append_history(req.message)
        elif req.status == STATUS_HUNTER:
            self.memory.append_history("Hunter/Wolf King is: " + req.name + ", they are activating their skill, choosing to shoot")
        elif req.status == STATUS_HUNTER_RESULT:
            if req.message:
                self.memory.append_history("Hunter/Wolf King is: " + req.name + ", they shot and took down " + req.message)
            else:
                self.memory.append_history("Hunter/Wolf King is: " + req.name + ", they didn't take anyone down")
        elif req.status == STATUS_SHERIFF_SPEECH_ORDER:
            if "Counter-clockwise" in req.message or "小号" in req.message:
                self.memory.append_history("Host: Sheriff speech order is small numbers first")
            else:
                self.memory.append_history("Host: Sheriff speech order is large numbers first")
        elif req.status == STATUS_SHERIFF_PK:
            self.memory.append_history(f"Sheriff PK speech: {req.name}: {req.message}")
        elif req.status == STATUS_RESULT:
            self.memory.append_history(req.message)
        else:
            raise NotImplementedError

    def interact(self, req=AgentReq) -> AgentResp:
        logger.info("wolf interact: {}".format(req))
        try:
            if req.status == STATUS_DISCUSS:
                if req.message:
                    self.memory.append_history(req.message)
                teammates = self.memory.load_variable("teammates")
                prompt = format_prompt(DESC_PROMPT,
                                      {"name": self.memory.load_variable("name"),
                                       "teammates": teammates,
                                       "history": "\n".join(self.memory.load_history())
                                      })
                logger.info("prompt:" + prompt)
                result = self.llm_caller(prompt)
                logger.info("wolf interact result: {}".format(result))
                return AgentResp(success=True, result=result, errMsg=None)

            elif req.status == STATUS_VOTE:
                self.memory.append_history('Host: It\'s time to vote. Everyone, please point to who you think might be a werewolf.')
                teammates = self.memory.load_variable("teammates")
                choices = [name for name in req.message.split(",")
                          if name != self.memory.load_variable("name") and name not in teammates]  # Exclude self and teammates
                self.memory.set_variable("choices", choices)
                prompt = format_prompt(VOTE_PROMPT, {"name": self.memory.load_variable("name"),
                                                   "teammates": teammates,
                                                   "choices": choices,
                                                   "history": "\n".join(self.memory.load_history())
                                                  })
                logger.info("prompt:" + prompt)
                result = self.llm_caller(prompt)
                logger.info("wolf interact result: {}".format(result))
                return AgentResp(success=True, result=result, errMsg=None)

            elif req.status == STATUS_WOLF_SPEECH:
                teammates = self.memory.load_variable("teammates")
                prompt = format_prompt(WOLF_SPEECH_PROMPT, {
                    "name": self.memory.load_variable("name"),
                    "teammates": teammates,
                    "history": "\n".join(self.memory.load_history())
                })
                logger.info("prompt:" + prompt)
                result = self.llm_caller(prompt)
                logger.info("wolf speech result: {}".format(result))
                return AgentResp(success=True, result=result, errMsg=None)

            elif req.status == STATUS_SKILL:
                teammates = self.memory.load_variable("teammates")
                choices = [name for name in req.message.split(",")
                          if name != self.memory.load_variable("name") and name not in teammates]  # Exclude self and teammates
                self.memory.set_variable("choices", choices)
                prompt = format_prompt(KILL_PROMPT, {
                    "name": self.memory.load_variable("name"),
                    "choices": choices,
                    "history": "\n".join(self.memory.load_history())
                })
                logger.info("prompt:" + prompt)
                result = self.llm_caller(prompt)
                logger.info("wolf kill result: {}".format(result))
                return AgentResp(success=True, result=result, skillTargetPlayer=result, errMsg=None)

            elif req.status == STATUS_SHERIFF_ELECTION:
                teammates = self.memory.load_variable("teammates")
                prompt = format_prompt(SHERIFF_ELECTION_PROMPT, {
                    "name": self.memory.load_variable("name"),
                    "teammates": teammates,
                    "history": "\n".join(self.memory.load_history())
                })
                logger.info("wolf agent sheriff election prompt:" + prompt)
                result = self.llm_caller(prompt)
                logger.info("wolf interact result: {}".format(result))
                return AgentResp(success=True, result=result, errMsg=None)

            elif req.status == STATUS_SHERIFF_SPEECH:
                teammates = self.memory.load_variable("teammates")
                prompt = format_prompt(SHERIFF_SPEECH_PROMPT, {
                    "name": self.memory.load_variable("name"),
                    "teammates": teammates,
                    "history": "\n".join(self.memory.load_history())
                })
                logger.info("wolf agent sheriff speech prompt:" + prompt)
                result = self.llm_caller(prompt)
                logger.info("wolf interact result: {}".format(result))
                return AgentResp(success=True, result=result, errMsg=None)

            elif req.status == STATUS_SHERIFF_PK:
                teammates = self.memory.load_variable("teammates")
                prompt = format_prompt(SHERIFF_SPEECH_PROMPT, {
                    "name": self.memory.load_variable("name"),
                    "teammates": teammates,
                    "history": "\n".join(self.memory.load_history())
                })
                logger.info("wolf agent sheriff pk prompt:" + prompt)
                result = self.llm_caller(prompt)
                logger.info("wolf interact result: {}".format(result))
                return AgentResp(success=True, result=result, errMsg=None)

            elif req.status == STATUS_SHERIFF_VOTE:
                teammates = self.memory.load_variable("teammates")
                choices = [name for name in req.message.split(",")]
                prompt = format_prompt(SHERIFF_VOTE_PROMPT, {
                    "name": self.memory.load_variable("name"),
                    "teammates": teammates,
                    "choices": choices,
                    "history": "\n".join(self.memory.load_history())
                })
                logger.info("wolf agent sheriff vote prompt:" + prompt)
                result = self.llm_caller(prompt)
                logger.info("wolf interact result: {}".format(result))
                return AgentResp(success=True, result=result, errMsg=None)

            elif req.status == STATUS_SHERIFF_SPEECH_ORDER:
                prompt = format_prompt(SHERIFF_SPEECH_ORDER_PROMPT, {
                    "name": self.memory.load_variable("name"),
                    "history": "\n".join(self.memory.load_history())
                })
                logger.info("wolf agent sheriff speech order prompt:" + prompt)
                result = self.llm_caller(prompt)
                logger.info("wolf interact result: {}".format(result))
                return AgentResp(success=True, result=result, errMsg=None)

            elif req.status == STATUS_SHERIFF:
                # Sheriff transfers badge
                teammates = self.memory.load_variable("teammates")
                choices = [name for name in req.message.split(",")
                          if name != self.memory.load_variable("name") and name not in teammates]
                prompt = format_prompt(SHERIFF_TRANSFER_PROMPT, {
                    "name": self.memory.load_variable("name"),
                    "teammates": teammates,
                    "choices": choices,
                    "history": "\n".join(self.memory.load_history())
                })
                logger.info("wolf agent sheriff transfer prompt:" + prompt)
                result = self.llm_caller(prompt)
                logger.info("wolf interact result: {}".format(result))
                return AgentResp(success=True, result=result, errMsg=None)

            return AgentResp(success=True, result=None, errMsg=None)
        except Exception as e:
            logger.error("WolfAgent interact failed", exc_info=True)
            return AgentResp(success=False, result=None, errMsg=str(e))