process_prompt.py

import os
import re
import random

class PromptGenerator():
    def __init__(self, prompt_dir="./prompts_reasoning"):

        prompt_file_name = ["Doctor", "Mafia", "Detective", "Vote_good", "Vote_bad", 
                            "Discussion_good", "Discussion_bad", "Discussion_Detective", "Deduction"]
        self.prompt_template = {}
        for i in range(len(prompt_file_name)):
            with open(f"{prompt_dir}/{prompt_file_name[i]}.md", "r", encoding='utf-8') as f:
                self.prompt_template[prompt_file_name[i]] = f.read()

        self.personalities = self.init_personalities()
        
        self.player_id = 0
        self.role = ""
        self.team = ""
        self.teammates = ""

        self.game_message = ""
        self.observation = [set() for _ in range(6)]
        self.past_statements = set()
        
        self.obs = []


        self.curr_phase = ""
        self.valid_target = ""
        self.characteristics = ""
        self.raw_player_statements = ""

    def process_observation(self, observation: str, characteristics_idx = -1):        
        phase = ""
        self.game_message = ""
        valid_targets = []
        prompt = None
        clear_discussion = False

        segments = re.split(r'(?=\[GAME\])|(?=\[Player \d+\])', observation.strip())
        segments = [seg.strip() for seg in segments if seg.strip()]

        for segment in segments:
            if segment.startswith("[Player"):
                match = re.match(r"\[Player (\d+)\]\s*(.*)", segment, re.DOTALL)
                if match:
                    player_id = int(match.group(1))
                    message = match.group(2).strip()
                    if player_id != self.player_id and message != "":
                        self.observation[player_id].add(message)
                        self.obs.append(f'Player {player_id} say: [{message.strip()}].')
                    else:
                        self.past_statements.add(f'You said: [{message.strip()}] last round.')
                    
                    self.raw_player_statements += f'Player {player_id} say: [{message.strip()}].\n'

            elif segment.startswith("[GAME]"):
                if "Welcome to Secret Mafia!" in segment:
                    num_players = re.search(r"You are Player (\d+).", segment)
                    role_match = re.search(r"Your role:\s*([^\n\r]+)", segment)
                    team_match = re.search(r"(?:Your team|Team):\s*([^\n\r]+)", segment)
                    teammates_match = re.search(r"Your teammates are:\s*(.+)", segment)

                    self.player_id = int(num_players.group(1)) if num_players else self.player_id
                    self.role = role_match.group(1).strip() if role_match else self.role
                    self.team = team_match.group(1).strip() if team_match else self.team
                    if teammates_match:
                        self.teammates = teammates_match.group(1).strip()

                    continue # Skip the welcome message

                elif "IS a Mafia member" in segment or "IS NOT a Mafia member" in segment or \
                     "was killed during the night" in segment or "was eliminated by vote" in segment or \
                     "No one was killed tonight." in segment:
                    self.game_message += f'{segment.strip()}\n'

                elif "Night has fallen. Mafia" in segment or "Night phase" in segment:
                    target_numbers = re.findall(r'\[(\d+)\]', segment)
                    valid_targets = [int(num) for num in target_numbers]
                    phase = "Night"

                elif "Voting phase" in segment:
                    target_numbers = re.findall(r'\[(\d+)\]', segment)
                    valid_targets = [int(num) for num in target_numbers]
                    phase = "Voting"

                elif "Day breaks." in segment:
                    phase = "Discussion"

        if phase == "Night":
            prompt = self.prompt_template[f'{self.role}']
            clear_discussion = True
        elif phase == "Voting":
            if "Mafia" in self.team or "Mafia" in self.role:
                prompt = self.prompt_template['Vote_bad']
            else:
                prompt = self.prompt_template['Vote_good']
        elif phase == "Discussion":
            if "Mafia" in self.team or "Mafia" in self.role:
                prompt = self.prompt_template['Discussion_bad']
            elif "Detective" in self.role:
                prompt = self.prompt_template['Discussion_Detective']
            else:
                prompt = self.prompt_template['Discussion_good']
        else:
            prompt = self.prompt_template['Discussion_good']

        self.characteristics = self.get_characteristics_prompt(characteristics_idx)

        prompt = prompt.replace('<|number|>', str(self.player_id))
        prompt = prompt.replace('<|role|>', self.role)
        prompt = prompt.replace('<|team|>', self.team)
        prompt = prompt.replace('<|teammates|>', self.teammates)
        prompt = prompt.replace('<|target|>', ', '.join([f'[Player {target}]' for target in valid_targets]))
        # prompt = prompt.replace('<|past_statements|>', "\n".join(sorted(self.past_statements)))
        prompt = prompt.replace('<|gameMSG|>', self.game_message)
        prompt = prompt.replace('<|characteristics|>', self.characteristics)

        self.curr_phase = phase
        self.valid_target = ', '.join([f'[Player {target}]' for target in valid_targets])

        return prompt, clear_discussion
    
    def process_deduction(self): 
        deduction_prompt_list = []
        for i in range(6):
            if len(self.observation[i]) == 0 or i == self.player_id:
                continue
            
            statements_list = sorted(self.observation[i])
            all_statements = "\n".join([f'Player {i} say: [{statement}]' for statement in statements_list])

            prompt = self.prompt_template['Deduction']
            prompt = prompt.replace('<|gameMSG|>', self.game_message)
            prompt = prompt.replace('<|statements|>', all_statements)
            prompt = prompt.replace('<|number|>', str(i))
            deduction_prompt_list.append(prompt)

        return deduction_prompt_list
    
    def get_characteristics_prompt(self, characteristics_idx = -1):
        if characteristics_idx == -1:
            characteristics_idx = random.randint(0, 3)
        elif characteristics_idx == 99:
            return ""

        return self.personalities[self.role][characteristics_idx]

    def get_system_prompt(self):
        s1 = "You are an expert in playing the social deduction game named Secret Mafia.\n"
        s2 = "The game has sin roles including two Mafia, one Detective, one Doctor, and two Villagers.\n"
        s3 = "There are six players including Player 0, Player 1, Player 2, Player 3, Player 4, Player 5.\n"
        s4 = "At the beginning of the game, each player is assigned a hidden role which divides them into the Mafia and the Villagers (Detective, Doctor, Villagers).\n"
        s5 = "Then the game alternates between the night round and the day round until one side wins the game.\n"
        s6 = "In the night round: the Mafia choose one player to kill; the Detective chooses one player to see if they are a Werewolf;"
        s7 = "the Doctor chooses one player including themselves to save without knowing who is chosen by the Mafia;"
        s8 = "the Villagers do nothing.\n"
        s9 = "In the day round: three round discussion phase, and a voting phase are performed in order.\n"
        s10 = "In the discussion phase, each remaining player speaks only once in order from Player 0 to Player 5 to discuss who might be the Mafia.\n"
        s11 = "In the voting phase, each player votes for one player or choose not to vote. The player with the most votes is eliminated and the game continues to the next night round.\n"
        s12 = "The Mafia win the game if the number of remaining Mafia is equal to the number of remaining Detective, Doctor, and Villagers.\n"
        s13 = "The Detective, Doctor, and Villagers win the game if all Mafia are eliminated."

        return s1 + s2 + s3 + s4 + s5 + s6 + s7 + s8 + s9 + s10 + s11 + s12 + s13
    
    def init_personalities(self):
        return {
            "Mafia": [
                "A calm and rational strategist. Uses logic to persuade others, analyzes every statement, and spots flaws.",
                "A passionate and outgoing leader. Charismatic, drives discussions, and inspires confidence.",
                "A cautious and conservative observer. Prefers to follow the group but gives key insights at critical moments.",
                "A witty and humorous negotiator. Eases tension with jokes while secretly gathering information."
            ],
            "Doctor": [
                "A kind and righteous protector. Actively joins discussions to defend the innocent and identify threats.",
                "A cautious and meticulous analyst. Quietly observes details and only speaks when certain.",
                "A gentle and friendly mediator. Avoids conflict, promotes harmony, and protects key people.",
                "A brave and straightforward justice defender. Boldly calls out suspicious behavior without fear."
            ],
            "Detective": [
                "A sharp and professional detective. Systematically analyzes clues and uncovers truth from details.",
                "A discreet and mysterious investigator. Blends in as a villager while secretly gathering intel.",
                "An intuitive and sensitive truth-seeker. Trusts instincts, notices emotions, and voices suspicions.",
                "A logical and rigorous reasoning expert. Builds deduction chains based only on logic and evidence."
            ],
            "Villager": [
                "A pure and kind villager. Trusts others and helps the group find villains to protect peace.",
                "A clever and quick-witted villager. Listens carefully and uses common sense to judge people.",
                "A cautious and suspicious villager. Keeps guard, questions everything, and doubts easily.",
                "A brave and righteous villager. Stands up for truth and justice, even in danger."
            ]
        }
    
    def get_round_info(self):
        round_info = {
            "role" : self.role.strip(),
            "team" : self.team.strip(),
            "teammates" : self.teammates.strip(),
            "characteristics" : self.characteristics.strip(),
            "phase" : self.curr_phase.strip(),
            "valid_target" : self.valid_target.strip(),
            "game_message" : self.game_message.strip(),
            "raw_player_statements" : self.raw_player_statements.strip()
        }

        return round_info

    def clear_observation(self):
        for i in range(6):
            self.observation[i].clear()
        self.obs.clear()
        self.raw_player_statements = ""

    def reset(self):
        self.player_id = 0
        self.role = ""
        self.team = ""
        self.teammates = ""

        self.game_message = ""
        self.observation = [set() for _ in range(6)]
        self.past_statements = set()
        
        self.obs = []

        self.curr_phase = ""
        self.valid_target = ""
        self.characteristics = ""
        self.raw_player_statements = ""