Upload hunter_agent.py

werewolf/hunter/hunter_agent.py

@@ -1,34 +1,192 @@
-import re
 from agent_build_sdk.model.roles import ROLE_HUNTER
-from agent_build_sdk.model.werewolf_model import AgentResp, AgentReq, STATUS_START, STATUS_SKILL, STATUS_DISCUSS, STATUS_VOTE, STATUS_HUNTER, STATUS_SHERIFF_ELECTION
+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
+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 hunter.prompt import GAME_RULE_PROMPT, DESC_PROMPT, SKILL_PROMPT, VOTE_PROMPT
+from hunter.prompt import DESC_PROMPT, VOTE_PROMPT, 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 HunterAgent(BasicRoleAgent):
+    """猎人角色Agent"""
+
     def __init__(self, model_name):
         super().__init__(ROLE_HUNTER, model_name=model_name)
+        self.memory.set_variable("can_shoot", True)  # 猎人初始可以开枪
 
-    def perceive(self, req: AgentReq):
+    def perceive(self, req=AgentReq):
         if req.status == STATUS_START:
             self.memory.clear()
+            self.memory.set_variable("name", req.name)
+            self.memory.set_variable("can_shoot", True)
             self.memory.append_history(GAME_RULE_PROMPT)
+            self.memory.append_history(f"主持人：你好，你分配到的角色是[猎人]，你是{req.name}")
+        elif req.status == STATUS_NIGHT:
+            self.memory.append_history("主持人：现在进入夜晚，天黑请闭眼")
+        elif req.status == STATUS_SKILL_RESULT:
+            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_RESULT:
+            self.memory.append_history(req.message)
+        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}")
+        else:
+            raise NotImplementedError
+
+    def interact(self, req=AgentReq) -> AgentResp:
+        logger.info("hunter interact: {}".format(req))
+        if req.status == STATUS_DISCUSS:
+            if req.message:
+                self.memory.append_history(req.message)
+            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"),
+                                   "shoot_info": shoot_info,
+                                   "history": "\n".join(self.memory.load_history())
+                                  })
+            logger.info("prompt:" + prompt)
+            result = self.llm_caller(prompt)
+            logger.info("hunter interact result: {}".format(result))
+            return AgentResp(success=True, result=result, errMsg=None)
 
-    def interact(self, req: AgentReq) -> AgentResp:
-        name = req.name if req.name else "我"
-        if req.status == STATUS_SHERIFF_ELECTION:
-            return AgentResp(success=True, result="否")
-            
-        hist = "\n".join(self.memory.load_history()[-10:])
-        if req.status in [STATUS_SKILL, STATUS_HUNTER]:
-            res = self.llm_caller(format_prompt(SKILL_PROMPT, {"history": hist, "name": name, "choices": req.message}))
-            target = re.search(r"(\d+)", res).group(1) if re.search(r"(\d+)", res) else ""
-            return AgentResp(success=True, result=res, skillTargetPlayer=target)
-        elif req.status == STATUS_DISCUSS:
-            res = self.llm_caller(format_prompt(DESC_PROMPT, {"history": hist, "name": name, "shoot_info": "可以开枪"}))
-            return AgentResp(success=True, result=res[:200])
         elif req.status == STATUS_VOTE:
-            res = self.llm_caller(format_prompt(VOTE_PROMPT, {"history": hist, "choices": req.message, "name": name}))
-            target = re.search(r"(\d+)", res).group(1) if re.search(r"(\d+)", res) else req.message.split(",")[0]
-            return AgentResp(success=True, result=target)
-        return AgentResp(success=True, result="猎人带队。")
+            self.memory.append_history('主持人: 到了投票的时候了。每个人，请指向你认为可能是狼人的人。')
+            choices = [name for name in req.message.split(",") if name != self.memory.load_variable("name")]
+            self.memory.set_variable("choices", choices)
+            prompt = format_prompt(VOTE_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("hunter interact result: {}".format(result))
+            return AgentResp(success=True, result=result, errMsg=None)
+
+        elif req.status == STATUS_SKILL:
+            # 猎人技能：开枪射杀一名玩家（遗言阶段）
+            can_shoot = self.memory.load_variable("can_shoot")
+            if not can_shoot:
+                return AgentResp(success=True, result="不开枪", errMsg=None)
+
+            choices = [name for name in req.message.split(",") if name != self.memory.load_variable("name")]
+            prompt = format_prompt(SKILL_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("hunter 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)
+
+        elif req.status == STATUS_SHERIFF_ELECTION:
+            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"),
+                                   "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:
+            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"),
+                                   "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:
+            choices = req.message.split(",")
+            prompt = format_prompt(SHERIFF_VOTE_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)
+            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:
+            # 警长转移警徽
+            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")]
+            prompt = format_prompt(SHERIFF_TRANSFER_PROMPT,
+                                  {"name": self.memory.load_variable("name"),
+                                   "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