Spaces:
Running
Running
| from agent_build_sdk.model.roles import ROLE_WITCH | |
| 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_SKILL | |
| 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 witch.prompt import DESC_PROMPT, VOTE_PROMPT, SKILL_PROMPT, GAME_RULE_PROMPT,CLEAN_USER_PROMPT | |
| class WitchAgent(BasicRoleAgent): | |
| """女巫角色Agent""" | |
| def __init__(self, model_name): | |
| super().__init__(ROLE_WITCH, model_name=model_name) | |
| # 初始化女巫的两瓶药 | |
| self.memory.set_variable("has_poison", True) | |
| self.memory.set_variable("has_antidote", 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("has_poison", True) | |
| self.memory.set_variable("has_antidote", True) | |
| self.memory.append_history(GAME_RULE_PROMPT) | |
| self.memory.append_history("主持人：你好，你分配到的角色是[女巫]") | |
| elif req.status == STATUS_NIGHT: | |
| self.memory.append_history("主持人：现在进入夜晚，天黑请闭眼") | |
| elif req.status == STATUS_SKILL_RESULT: | |
| self.memory.append_history(f"主持人：女巫，你使用技能的结果是{req.message}") | |
| 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('主持人: 每个玩家描述自己的信息。') | |
| elif req.status == STATUS_VOTE: # 投票环节 | |
| self.memory.append_history(f'第{req.round}天的投票环节,{req.name} 投了 {req.message}') | |
| elif req.status == STATUS_VOTE_RESULT: # 投票结果 | |
| out_player = req.name if req.name else req.message | |
| if out_player: | |
| self.memory.append_history('主持人: 投票结果是：{}。'.format(out_player)) | |
| else: | |
| self.memory.append_history('主持人: 无人出局。') | |
| elif req.status == STATUS_RESULT: | |
| self.memory.append_history(req.message) | |
| else: | |
| raise NotImplementedError | |
| def interact(self, req=AgentReq) -> AgentResp: | |
| logger.info("witch interact: {}".format(req)) | |
| if req.status == STATUS_DISCUSS: | |
| if req.message: | |
| self.memory.append_history(req.message) | |
| has_poison = self.memory.load_variable("has_poison") | |
| has_antidote = self.memory.load_variable("has_antidote") | |
| skill_info = "女巫有{}瓶毒药和{}瓶解药".format("1" if has_poison else "0", "1" if has_antidote else "0") | |
| prompt = format_prompt(DESC_PROMPT, | |
| {"name": self.memory.load_variable("name"), | |
| "skill_info": skill_info, | |
| "history": "\n".join(self.memory.load_history()) | |
| }) | |
| logger.info("prompt:" + prompt) | |
| result = self.llm_caller(prompt) | |
| logger.info("witch interact result: {}".format(result)) | |
| return AgentResp(success=True, result=result, errMsg=None) | |
| elif req.status == STATUS_VOTE: | |
| 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("witch interact result: {}".format(result)) | |
| return AgentResp(success=True, result=result, errMsg=None) | |
| elif req.status == STATUS_SKILL: | |
| has_poison = self.memory.load_variable("has_poison") | |
| has_antidote = self.memory.load_variable("has_antidote") | |
| tonight_killed = req.message | |
| skill_info = "女巫有{}瓶毒药和{}瓶解药".format("1" if has_poison else "0", "1" if has_antidote else "0") | |
| prompt = format_prompt(SKILL_PROMPT, { | |
| "name": self.memory.load_variable("name"), | |
| "tonight_killed": tonight_killed, | |
| "skill_info": skill_info, | |
| "history": "\n".join(self.memory.load_history()) | |
| }) | |
| logger.info("prompt:" + prompt) | |
| result = self.llm_caller(prompt) | |
| logger.info("witch skill result: {}".format(result)) | |
| # 根据结果更新药水状态 | |
| skill_target_person = None | |
| if result.startswith("救") and has_antidote: | |
| self.memory.set_variable("has_antidote", False) | |
| self.memory.append_history(f"女巫使用解药救活了{tonight_killed}") | |
| skill_target_person = tonight_killed | |
| elif result.startswith("毒") and has_poison: | |
| poisoned_player = result[1:].strip() | |
| self.memory.set_variable("has_poison", False) | |
| self.memory.append_history(f"女巫使用毒药杀死了{poisoned_player}") | |
| skill_target_person = poisoned_player | |
| return AgentResp(success=True, result=result, skillTargetPlayer=skill_target_person, errMsg=None) | |
| else: | |
| raise NotImplementedError | |