werewolf/wolf_king/wolf_king_agent.py

from agent_build_sdk.model.roles import ROLE_WOLF_KING
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_VOTE, STATUS_SHERIFF, \
    STATUS_SHERIFF_SPEECH_ORDER, STATUS_SHERIFF_PK, 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_king.prompt import DESC_PROMPT, VOTE_PROMPT, WOLF_SPEECH_PROMPT, KILL_PROMPT, SHOOT_SKILL_PROMPT, \
    GAME_RULE_PROMPT, CLEAN_USER_PROMPT, SHERIFF_ELECTION_PROMPT, SHERIFF_SPEECH_PROMPT, SHERIFF_VOTE_PROMPT, \
    SHERIFF_SPEECH_ORDER_PROMPT, SHERIFF_TRANSFER_PROMPT


class WolfKingAgent(BasicRoleAgent):
    """狼王角色Agent"""

    def __init__(self, model_name):
        super().__init__(ROLE_WOLF_KING, model_name=model_name)
        self.memory.set_variable("teammates", [])
        self.memory.set_variable("can_shoot", True)  # 狼王可以开枪

    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", [])
            self.memory.set_variable("can_shoot", True)
            self.memory.append_history(GAME_RULE_PROMPT)
            self.memory.append_history(f"主持人：你好，你分配到的角色是[狼王]，你是{req.name}")
            if req.message:  # 如果有队友信息
                teammates = req.message.split(",")
                self.memory.set_variable("teammates", teammates)
                self.memory.append_history(f"主持人：你的狼人队友是: {req.message}")
        elif req.status == STATUS_NIGHT:
            self.memory.append_history("主持人：现在进入夜晚，天黑请闭眼")
        elif req.status == STATUS_WOLF_SPEECH:
            # 狼人之间的交流
            if req.name:
                self.memory.append_history(f"狼人{req.name}说: {req.message}")
            else:
                self.memory.append_history("主持人：狼人请睁眼，狼人请互相确认身份，并选择要击杀的对象")
        elif req.status == STATUS_SKILL_RESULT:
            if req.name:
                # 如果是击杀结果
                self.memory.append_history(f"主持人：狼人今晚选择击杀的目标是:{req.name}")
            elif req.message:
                # 如果是开枪结果
                self.memory.append_history(f"主持人：{req.message}")
                if "能开枪" in req.message:
                    self.memory.set_variable("can_shoot", True)
                elif "不能开枪" in req.message:
                    self.memory.set_variable("can_shoot", False)
        elif req.status == STATUS_NIGHT_INFO:
            self.memory.append_history(f"主持人：天亮了！昨天晚上的信息是: {req.message}")
        elif req.status == STATUS_DISCUSS:  # 发言环节
            if req.name:
                # 其他玩家发言
                # 可以使用模型来过滤掉玩家的注入消息，也可以换一个小模型，实际使用需要考虑对memory加锁，避免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:
                # 主持人发言
                self.memory.append_history('主持人: 现在进入第{}天。'.format(str(req.round)))
                self.memory.append_history('主持人: 每个玩家描述自己的信息。')
            self.memory.append_history("---------------------------------------------")
        elif req.status == STATUS_VOTE:  # 投票环节
            self.memory.append_history(f'第{req.round}天。投票信息:{req.name}投了{req.message}')
        elif req.status == STATUS_VOTE_RESULT:  # 投票环节
            if req.name:
                self.memory.append_history('主持人: 投票结果是：{}。'.format(req.name))
            else:
                self.memory.append_history('主持人: 无人出局。')
        elif req.status == STATUS_SHERIFF_ELECTION:
            self.memory.append_history(f"主持人: 上警玩家: {req.message}")
        elif req.status == STATUS_SHERIFF_SPEECH:
            self.memory.append_history(f"{req.name} (警上发言): {req.message}")
        elif req.status == STATUS_SHERIFF_VOTE:
            self.memory.append_history(f"警上投票: {req.name}投了{req.message}")
        elif req.status == STATUS_SHERIFF:
            if req.name:
                self.memory.append_history(f"主持人: 警徽归属: {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("猎人/狼王是：" + req.name + ",他正在发动技能，选择开枪")
        elif req.status == STATUS_HUNTER_RESULT:
            if req.message:
                self.memory.append_history("猎人/狼王是：" + req.name + ",他开枪带走了" + req.message)
            else:
                self.memory.append_history("猎人/狼王是：" + req.name + ",他没有带走任何人")
        elif req.status == STATUS_SHERIFF_SPEECH_ORDER:
            if "小号" in req.message:
                self.memory.append_history("主持人: 警长发言顺序是小号优先")
            else:
                self.memory.append_history("主持人: 警长发言顺序是大号优先")
        elif req.status == STATUS_SHERIFF_PK:
            self.memory.append_history(f"警长PK发言: {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 king interact: {}".format(req))
        if req.status == STATUS_DISCUSS:
            if req.message:
                self.memory.append_history(req.message)
            teammates = self.memory.load_variable("teammates")
            can_shoot = self.memory.load_variable("can_shoot")
            shoot_info = "可以开枪" if can_shoot else "已失去开枪能力"
            prompt = format_prompt(DESC_PROMPT,
                                  {"name": self.memory.load_variable("name"),
                                   "teammates": teammates,
                                   "shoot_info": shoot_info,
                                   "history": "\n".join(self.memory.load_history())
                                  })
            logger.info("prompt:" + prompt)
            result = self.llm_caller(prompt)
            logger.info("wolf king interact result: {}".format(result))
            return AgentResp(success=True, result=result, errMsg=None)

        elif req.status == STATUS_VOTE:
            self.memory.append_history('主持人: 到了投票的时候了。每个人，请指向你认为可能是狼人的人。')
            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]  # 排除自己和队友
            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 king 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 king speech result: {}".format(result))
            return AgentResp(success=True, result=result, errMsg=None)

        elif req.status == STATUS_SKILL:
            # 判断是击杀技能还是开枪技能
            message = req.message
            if message and "请发表最后的遗言" in message:
                # 开枪技能：狼王被淘汰时的开枪
                can_shoot = self.memory.load_variable("can_shoot")
                if not can_shoot:
                    return AgentResp(success=True, result="不开枪", errMsg=None)

                teammates = self.memory.load_variable("teammates")
                choices = [name for name in message.replace("请发表最后的遗言", "").split(",")
                          if name and name.strip() and name != self.memory.load_variable("name") and name not in teammates]

                prompt = format_prompt(SHOOT_SKILL_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 king shoot skill result: {}".format(result))

                if result != "不开枪":
                    self.memory.set_variable("can_shoot", False)

                return AgentResp(success=True, result=result, skillTargetPlayer=None if result == "不开枪" else result, errMsg=None)
            else:
                # 击杀技能：狼人夜晚击杀
                teammates = self.memory.load_variable("teammates")
                choices = [name for name in message.split(",")
                          if name != self.memory.load_variable("name") and name not in teammates]
                self.memory.set_variable("choices", choices)

                prompt = format_prompt(KILL_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 king 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")
            can_shoot = self.memory.load_variable("can_shoot")
            shoot_info = "可以开枪" if can_shoot else "已失去开枪能力"
            prompt = format_prompt(SHERIFF_ELECTION_PROMPT,
                                  {"name": self.memory.load_variable("name"),
                                   "teammates": teammates,
                                   "shoot_info": shoot_info,
                                   "history": "\n".join(self.memory.load_history())
                                  })
            logger.info("prompt:" + prompt)
            result = self.llm_caller(prompt)
            return AgentResp(success=True, result=result, errMsg=None)

        elif req.status == STATUS_SHERIFF_SPEECH or req.status == STATUS_SHERIFF_PK:
            teammates = self.memory.load_variable("teammates")
            can_shoot = self.memory.load_variable("can_shoot")
            shoot_info = "可以开枪" if can_shoot else "已失去开枪能力"
            prompt = format_prompt(SHERIFF_SPEECH_PROMPT,
                                  {"name": self.memory.load_variable("name"),
                                   "teammates": teammates,
                                   "shoot_info": shoot_info,
                                   "history": "\n".join(self.memory.load_history())
                                  })
            logger.info("prompt:" + prompt)
            result = self.llm_caller(prompt)
            return AgentResp(success=True, result=result, errMsg=None)

        elif req.status == STATUS_SHERIFF_VOTE:
            teammates = self.memory.load_variable("teammates")
            choices = 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("prompt:" + prompt)
            result = self.llm_caller(prompt)
            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("prompt:" + prompt)
            result = self.llm_caller(prompt)
            return AgentResp(success=True, result=result, errMsg=None)

        elif req.status == STATUS_SHERIFF:
            # 警长转移警徽
            teammates = self.memory.load_variable("teammates")
            can_shoot = self.memory.load_variable("can_shoot")
            shoot_info = "可以开枪" if can_shoot else "已失去开枪能力"
            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,
                                   "shoot_info": shoot_info,
                                   "choices": choices,
                                   "history": "\n".join(self.memory.load_history())
                                  })
            logger.info("prompt:" + prompt)
            result = self.llm_caller(prompt)
            return AgentResp(success=True, result=result, errMsg=None)
        else:
            raise NotImplementedError