main.py

import argparse
import json
import os
import glob
from utils.Chat_utils import ChatSession
from utils.RAG import VectorDatabase
from utils.llama import load_llama_model
from utils.gemma import load_gemma_model
from utils.json_tools import json_format
import yaml
import pandas as pd
import numpy as np
import math
from tqdm import tqdm
from datetime import datetime
import pickle
class ChatItem:
    _index_counter = 0
    def __init__(self, stage='', speaker='', listener='', private=False, words='', prompt='', is_ask=1):
        self.index = ChatItem._index_counter
        ChatItem._index_counter += 1
        self.stage = stage  # the stage of the chat
        self.speaker = speaker  # the speaker of the chat
        self.listener = listener  # the listener of the chat
        self.private = private  # whether the chat is private
        self.words = words  # the content of the chat
        self.is_ask = is_ask  # the prompt of the chat
        self.prompt = prompt  # the prompt of the chat


class EvaluateItem:
    _index_counter = 0
    def __init__(self, tester='', question='', answer='', choices='', is_right='', truth='', value=''):
        self.index = ChatItem._index_counter
        ChatItem._index_counter += 1
        self.tester = tester  # the tester of the question
        self.question = question
        self.choices = choices
        self.truth = truth
        self.answer = answer
        self.is_right = is_right
        self.value = value


class Agent:

    def __init__(self, args, name = '', script='', acts_goal='', final_goal='', secrets='', acts_performance="", is_murderer=False, question=[], model=None, tokenizer=None):
        self.session = ChatSession(args=args, model=model, tokenizer=tokenizer)  # get reply
        self.name = name
        self.script = script  # current script
        self.acts_goal = acts_goal
        self.final_goal = final_goal
        self.questions = question
        self.acts_performance = acts_performance
        self.secrets = secrets
        self.is_murderer = is_murderer  # murderer or civilian
        self.database = VectorDatabase()
        self.summary = ''
        self.character_suspect = []

    def speak(self, prompt, evaluate = False):
        reply = self.session.get_reply(prompt, evaluate)
        return reply

class ScriptManager:

    def __init__(self, args):
        # {character：script(include goal, script, performance...)}
        self.agents_scripts = {}
        # {character：evaluate questions}
        # self.agents_evaluate = {}
        self.agents_final_result = {}
        # script parameters
        para_file_path = os.path.join(args.script_root, args.script_name, 'json', 'script_info.json')
        para = self.read_json(para_file_path)
        self.agent_num = para["agent_num"]  # script character number
        self.character_name_list = para["character_name"]  # script character name
        self.name2index = {name:str(i) for i, name in enumerate(self.character_name_list)}
        self.index2name = {str(i):name for i, name in enumerate(self.character_name_list)}
        self.script_name = para["script_name"]  # script  name
        self.open_discuss_rounds = para["open_discuss_rounds"]  # open talk rounds
        self.acts_num = para["acts_num"]  # acts number
        self.keep_asking = para["keep_asking"]  # continue asking number
        self.privacy_discuss_rounds = para["privacy_discuss_rounds"]  # privacy talk rounds
        # read each character's script
        for name in self.character_name_list:
            script_file_path = os.path.join(args.script_root, args.script_name, 'json', name + '.json')
            script_data = self.read_json(script_file_path)
            self.agents_scripts[name] = script_data
        for name in self.character_name_list:
            final_result_file_path = os.path.join(args.script_root, args.script_name, 'final_result', name + '.csv')
            final_result = pd.read_csv(final_result_file_path, encoding='utf-8-sig')
            self.agents_final_result[name] = final_result
        self.agents_final_result['GM'] = pd.read_csv(os.path.join(args.script_root, args.script_name, 'final_result', 'FSA.csv'), encoding='utf-8-sig')

    def update_script(self, agent, n):
        """
        update agent's script according to the acts number
        """
        if n > self.agent_num - 1:
            print('error')
            return
        name = agent.name
        agent.is_murderer = True if self.agents_scripts[name]["is_murderer"] == 1 else False
        agent.character_name_list = [i for i in self.character_name_list if i != name]

        agent.script = '\n'.join(self.agents_scripts[name]["script"][:n + 1])
        agent.acts_goal = '\n'.join(self.agents_scripts[name]["acts_goal"][:n + 1])
        # agent.acts_performance = self.agents_scripts[name]["acts_performance"][n]
        # agent.secrets = self.agents_scripts[name]["secrets"]
        agent.victims = self.agents_scripts[name]["victims"]
        agent.kill_by_me = self.agents_scripts[name]["kill_by_me"]
        # agent.questions = self.agents_scripts[name]["questions"][n]
        if n == self.agent_num - 1:
            agent.final_goal = self.agents_scripts[name]["final_goal"]
        if agent.character_suspect == []:
            agent.character_suspect = [[i for i in self.character_name_list if i != name]]*len(agent.victims)
        return agent


    def read_json(self, path):
        with open(path, 'r', encoding='utf-8') as file:
            data = json.load(file)
        return data

    def get_script(self, agent_name):
        return self.agents_scripts.get(agent_name, {})


class Game:
    def __init__(self, scriptmanager, args):
        # read prompts
        self.args = args
        with open(args.prompt_path, 'r', encoding='utf-8') as f:
            self.prompts = yaml.load(f, Loader=yaml.FullLoader)
            f.close()
        with open(args.sensor_path, 'r', encoding='utf-8') as f:
            self.sensors = json.load(f)
            f.close()
    # the number of characters who can ask questions when other characters introduce themselves
        self.intro_question_num = 2
        self.script_manager = scriptmanager
        self.agents = [] # store all the agents
        self.GM_agent = Agent(args=self.args) # store all the script
        self.current_stage = None  # current stage
        self.chat_history = []  # store all the chat history
        self.chat_dataset = VectorDatabase()
        self.evaluate_history = []

        self.rsl = args.rsl  # the max length for relative script length (retrival from current script by RAG)
        self.each_rsl = args.each_rsl  # the max length for relative script length (retrival from current script by RAG)
        self.rchl = args.rchl  # the max length for relative chat history length (retrival from chat history by RAG)
        self.question_number = args.question_number

        if "llama" in args.model_type:
            self.model, self.tokenizer = load_llama_model(args.ckpt_dir)
        elif "gemma" in args.model_type:
            self.model, self.tokenizer = load_gemma_model(args.ckpt_dir)

        self.initial_game()  # initial game

    def initial_game(self):
        """
        initial game
        """
        character_name_list = self.script_manager.character_name_list
        # initial agents
        for character_name in character_name_list:
            if "llama" in self.args.model_type or "gemma" in self.args.model_type:
                self.agents.append(Agent(args=self.args, name=character_name, model=self.model, tokenizer=self.tokenizer))
            else:
                self.agents.append(Agent(args=self.args, name=character_name))
        self.update_script(0)  # initial agent script(get act 0 script)
        self.initial_agent_session()  # initial agent system prompt

    def add_chat(self, chatitem):
        """
            add chat item to chat history and chat dataset (for RAG)
        """
        self.chat_history.append(chatitem)  # add chat item to chat history
        format_chat_item = self.format_chat_item(chatitem)  # get chat string for RAG
        # self.chat_dataset.add_text(format_chat_item, False)  # store chat string to chat dataset
        print(format_chat_item)
    def add_chat_database(self, last_n):
        chat = ''
        for i in range(last_n, 0, -1):
            chat += self.format_chat_item(self.chat_history[-i]) + '\n'  # get chat string for RAG
        self.chat_dataset.add_text(chat, False)  # store chat string to chat dataset

    def format_chat_item(self, chatitem):
        """
        convert chat item to a string that can be used by RAG
        """
        if chatitem.is_ask == 1:
            format_chat_item = self.prompts['combine_chat_ask'].format(speaker=chatitem.speaker,listener=chatitem.listener, words=chatitem.words)
        elif chatitem.is_ask == 2:
            format_chat_item = self.prompts['combine_chat_all'].format(speaker=chatitem.speaker,listener=chatitem.listener, words=chatitem.words)
        else:
            format_chat_item = self.prompts['combine_chat_answer'].format(speaker=chatitem.speaker,listener=chatitem.listener, words=chatitem.words)

        return format_chat_item


    def get_chat_history(self, name):
        """
        This function has been deprecated.
        """
        chat_history_list = []
        for chatitem in self.chat_history:
            if chatitem.private == True and chatitem.listener == name:
                chat_history_list.append(self.format_chat_item(chatitem))
            if chatitem.private == False:
                chat_history_list.append(self.format_chat_item(chatitem))
        return ', '.join(chat_history_list)

    def update_script(self, n):
        """
        update agent's script according to the acts number
        add the updated script to the script dataset
        """
        all_scrpt = ''
        for i in range(len(self.agents)):
            self.agents[i] = self.script_manager.update_script(self.agents[i], n)
            all_scrpt += self.agents[i].script
            file_dir = os.path.join(args.database_root, args.script_name.split("-")[0])
            file_name_faiss = os.path.join(args.database_root, args.script_name.split("-")[0],  ''.join([self.script_manager.name2index[self.agents[i].name],'_act_',str(n),'.faiss']))
            file_name_doc = os.path.join(args.database_root, args.script_name.split("-")[0],  ''.join([self.script_manager.name2index[self.agents[i].name],'_act_',str(n),'.pickle']))
            if os.path.exists(file_name_faiss):
                self.agents[i].database.load_faiss(file_name_faiss)
                self.agents[i].database.load_texts(file_name_doc)

            else:
                self.agents[i].database.reset()
                self.agents[i].database.add_text(self.agents[i].script)
                if not os.path.exists(file_dir):
                    os.makedirs(file_dir)
                self.agents[i].database.save_faiss(file_name_faiss)
                self.agents[i].database.save_texts(file_name_doc)


    def initial_agent_session(self):
        """initial agent system prompt according to the role (civilian or murderer)"""
        self.GM_agent.session.system = self.prompts['system_prompt_gm']
        for i in range(len(self.agents)):
            if self.agents[i].is_murderer:
                victim = ",".join([victim for j, victim in enumerate(self.agents[i].victims) if self.agents[i].kill_by_me[j] == 1])
                self.agents[i].session.system = self.prompts['system_prompt_murderer'].format(
                    character_name=self.agents[i].name,
                    victims = victim,
                    character_name_list=', '.join(self.agents[i].character_name_list))
            else:
                self.agents[i].session.system = self.prompts['system_prompt_civilian'].format(
                    character_name=self.agents[i].name,
                    character_name_list=', '.join(self.agents[i].character_name_list))

    def get_agent(self,name):
        """ get agent by character name"""
        return [i for i in self.agents if i.name == name][0]

    def self_introduction_phase(self):
        stage = 'self introduction stage'
        for i in tqdm(range(len(self.agents))):
            # Add host dialogue
            # self.add_chat(ChatItem(stage, "hold", self.agents[i].name, False,self.prompts['hold_intro_yourself'].format(character=self.agents[i].name)))
            if not args.is_english:

                current_script = self.agents[i].database.query('你是谁？', 12000)
            else:
                current_script = self.agents[i].database.query('who are you?', 12000)
            if self.agents[i].is_murderer:
                intro_prompt = self.prompts['self_intro_intro'].format(current_script=current_script,   goal=self.agents[i].acts_goal)
            else:
                intro_prompt = self.prompts['self_intro_intro_murder'].format(current_script=current_script,   goal=self.agents[i].acts_goal)
            intro = self.agents[i].speak(intro_prompt)
            # intro_pro = self.process_reply(intro)
            self.add_chat(ChatItem(stage, self.agents[i].name, "all", False, intro,intro_prompt, is_ask=2))
            self.add_chat_database(1)

    def find_most_suspicion(self, i, x, victim, murderer = False, multi = True):
        postfix = ''
        if murderer:
            postfix = '_murder'
        summary_list = []
        for j in range(len(self.agents)):
            if i == j:
                continue
            if self.agents[j].name not in self.agents[i].character_suspect[x]:
                continue
            current_script = self.agents[i].database.query(self.agents[j].name+','+victim, self.rsl)
            chat_history = self.chat_dataset.query(self.agents[j].name+','+victim, self.rsl)
            summary_single = []
            for i in range(len(self.sensors)):
                sensor_single = self.sensors[str(i)]
                choices = sensor_single["choices"]
                choices = ' or '.join(['"'+ item + '"' for item in sensor_single["choices"]])
                sensor_prompt = self.prompts['sensor_ask'+postfix].replace('{current_script}',current_script).replace('{chat_history}',chat_history).replace('{character}',self.agents[j].name).replace('{victim}',victim).replace('{sensor}',sensor_single["prompt"]).replace('{choices}',choices)
                sensor_reply = self.agents[i].speak(sensor_prompt)
                sensor_summary_single = self.prompts["sensor_summary_single"].format(name=sensor_single["summary"],character=self.agents[j].name,opinion=sensor_reply)
                summary_single.append(sensor_summary_single)
            summary_str = "About {character}:\n {summary}".format(character = self.agents[j].name, summary = '\n'.join(summary_single))
            summary_list.append(summary_str)

        summary = '\n'.join(summary_list)
        character_suspect = str(self.agents[i].character_suspect[x])
        number = math.ceil(len(self.agents[i].character_suspect[x]) / 2)

        # if self.args.mode:
        ask_question_prompt = self.prompts['s_constraint_multi'+postfix].replace('{character}',self.agents[j].name).replace('{summary}',summary).replace('{character_suspect}',character_suspect).replace('{victim}',victim)
        # else:
        #     ask_question_prompt = self.prompts['s_constraint'+postfix].replace('{character}',self.agents[j].name).replace('{summary}',summary).replace('{number}',str(number)).replace('{character_suspect}',character_suspect).replace('{victim}',victim)
        update_suspect = self.agents[i].speak(ask_question_prompt)
        format_update_suspect = json_format(update_suspect)
        return format_update_suspect, ask_question_prompt

    def constraints_space(self):
        for i in tqdm(range(len(self.agents)), desc='constraints_space'):
            for x,victim in enumerate(self.agents[i].victims):
                # get question
                if self.agents[i].kill_by_me[x] == 1:
                    format_update_suspect, ask_question_prompt = self.find_most_suspicion(i, x, victim, True)
                else:
                    format_update_suspect, ask_question_prompt = self.find_most_suspicion(i, x, victim, False)
                try:
                    self.agents[i].character_suspect[x] = format_update_suspect['suspicion']
                except:
                    continue

    def propose_question_talk_stage(self, i, j, victim, murderer = False):
        postfix = ''
        if murderer:
            postfix = '_murder'
        if self.args.sensor:
            current_script = self.agents[i].database.query(self.agents[j].name+','+victim, self.rsl)
            chat_history = self.chat_dataset.query(self.agents[j].name+','+victim, self.rsl)
            summary = []
            for i in range(len(self.sensors)):
                sensor_single = self.sensors[str(i)]
                if sensor_single['question_ask'] == 1:
                    # choices = sensor_single["choices"]
                    choices = ' or '.join(['"'+ item + '"' for item in sensor_single["choices"]])
                    sensor_prompt = self.prompts['sensor_ask'+postfix].replace('{current_script}',current_script).replace('{chat_history}',chat_history).replace('{character}',self.agents[j].name).replace('{victim}',victim).replace('{sensor}',sensor_single["prompt"]).replace('{choices}',choices)
                    sensor_reply = self.agents[i].speak(sensor_prompt)
                    sensor_summary_single = self.prompts["sensor_summary_single"].format(name=sensor_single["summary"],character=self.agents[j].name,opinion=sensor_reply)
                    summary.append(sensor_summary_single)
            summary_str = '\n'.join(summary)
        else:
            summary_str = ''
        current_script = self.agents[i].database.query(self.agents[j].name+','+victim, self.rsl)
        chat_history = self.chat_dataset.query(self.agents[j].name+','+victim, self.rsl)
        ask_question_prompt = self.prompts['sensor_question_get'+postfix].replace('{character}',self.agents[j].name).replace('{summary}',summary_str).replace('{victim}',victim).replace('{question_number}',self.question_number).replace('{current_script}',current_script).replace('{chat_history}',chat_history)
        questions = self.agents[i].speak(ask_question_prompt)
        format_question = json_format(questions)
        return format_question, ask_question_prompt

    def open_talk(self):
        for i in tqdm(range(len(self.agents))):
            for x,victim in enumerate(self.agents[i].victims):
                # get question
                for j in tqdm(range(len(self.agents)), desc='Character {0} investigate {1}'.format(self.agents[i].name, victim)):
                    stage = '{0} ask {1} for death of {2}'.format(self.agents[i].name, self.agents[j].name, victim)
                    if i == j:
                        continue
                    if self.agents[j].name not in self.agents[i].character_suspect[x]:
                        continue
                    # question_, question_prompt = self.propose_question_talk_stage(i)

                    if self.agents[i].kill_by_me[x] == 1:
                        questions, question_prompt = self.propose_question_talk_stage(i, j, victim, True)
                    else:
                        questions, question_prompt = self.propose_question_talk_stage(i, j, victim, False)

                    try:
                        for key, question in questions.items():
                            self.add_chat(ChatItem(stage, self.agents[i].name, self.agents[j].name, False, question, question_prompt, is_ask=1))
                            # get relay agent
                            relay_agent = self.agents[j]
                            # reply:
                            current_script = relay_agent.database.query(question, self.rsl)
                            chat_history = self.chat_dataset.query(question, self.rchl)

                            if relay_agent.kill_by_me[x] == 1:
                                reply_prompt = self.prompts['open_talk_reply_murder'].format(current_script=current_script,goal=relay_agent.acts_goal,chat_history=chat_history,character=self.agents[i].name,question=question, victim=victim)
                            else:
                                reply_prompt = self.prompts['open_talk_reply'].format(current_script=current_script,goal=relay_agent.acts_goal,chat_history=chat_history,character=self.agents[i].name,question=question)

                            # reply_prompt = self.prompts['open_talk_reply_rap'].format(current_script=current_script,goal=relay_agent.acts_goal,chat_history=chat_history,character=self.agents[i].name,question=question)
                            reply = relay_agent.speak(reply_prompt)
                            reply_pro = self.process_reply(reply)
                            # if reply!=reply_pro:
                            #     print(reply)
                            self.add_chat(ChatItem(stage, relay_agent.name, self.agents[i].name, False, reply_pro, reply_prompt, is_ask=0))
                            self.add_chat_database(2)
                    except:
                        continue

    def process_reply(self, parts):
        # Split the string based on the first occurrence of ":" or "："
        # parts = reply.split(':')[-1] if ':' in reply else reply
        # parts = parts.split('：')[-1] if '：' in parts else parts

        # Remove text within the last set of parentheses
        if '(' in parts and ')' in parts:
            start = parts.rfind('(')
            end = parts.rfind(')')
            parts = parts[:start] + parts[end+1:]
        elif '（' in parts and '）' in parts:
            start = parts.rfind('（')
            end = parts.rfind('）')
            parts = parts[:start] + parts[end+1:]

        return parts

    def save_history(self,path,prefix = ''):
        history_list = []
        for chatitem in self.chat_history:
            history_list.append({
                "index" : chatitem.index,
                "stage" : chatitem.stage,
                "speaker" : chatitem.speaker,
                "listener" : chatitem.listener,
                # "private" : chatitem.private,
                "words" : chatitem.words,
                "prompt" : chatitem.prompt
            })
        history = pd.DataFrame(history_list)

        path_excel = os.path.join(path,prefix+self.script_manager.script_name+'.xlsx')
        history.to_excel(path_excel, index=False,  encoding='utf-8-sig')
        print('save chat history to excel')

        file_name_faiss = os.path.join(path, prefix+self.script_manager.script_name+'chat_history.faiss')
        file_name_doc = os.path.join(path, prefix+self.script_manager.script_name+'chat_history.pickle')
        self.chat_dataset.save_faiss(file_name_faiss)
        self.chat_dataset.save_texts(file_name_doc)

    def load_history(self, path):
        """
        load chat history from excel
        """
        # ChatItem._index_counter = 0
        # path_excel = os.path.join(path, self.script_manager.script_name + '.xlsx')
        # history_df = pd.read_excel(path_excel)
        # self.chat_history = []
        # for _, row in history_df.iterrows():
        #     chat_item = ChatItem(stage=row['stage'],
        #                          speaker=row['speaker'],
        #                          listener=row['listener'],
        #                          # private=row['private'],
        #                          words=row['words'],
        #                          prompt=row['prompt'])
        #     chat_item.index = row['index']
        #     self.chat_history.append(chat_item)
        #     format_chat_item = self.format_chat_item(chat_item)
        #     self.chat_dataset.add_text(format_chat_item, False)
        # if self.chat_history:
        #     ChatItem._index_counter = self.chat_history[-1].index + 1
        # self.update_script(0)
        file_name_faiss = os.path.join(path, self.script_manager.script_name+'chat_history.faiss')
        file_name_doc = os.path.join(path, self.script_manager.script_name+'chat_history.pickle')

        self.chat_dataset.load_faiss(file_name_faiss)
        self.chat_dataset.load_texts(file_name_doc)



    def save_evaluate_history(self,path,prefix = ''):
        path = os.path.join(path, 'evaluate')
        if not os.path.exists(path):
            os.makedirs(path)
        history_list = []
        for chatitem in self.evaluate_history:
            if chatitem.value == "a":
                chatitem.value = 10
            elif chatitem.value == "b":
                chatitem.value = 5
            elif chatitem.value == "c":
                chatitem.value = 2
            history_list.append({
                "tester" : chatitem.tester,
                "value" : chatitem.value,
                "question" : chatitem.question,
                "choices" : chatitem.choices.replace('\n',' '),
                "truth" : chatitem.truth,
                "answer" : chatitem.answer,
                "is_right" : chatitem.is_right

            })
        history = pd.DataFrame(history_list)
        type_a = history[history['value'] == 10]['is_right'].mean()
        type_b = history[history['value'] == 5]['is_right'].mean()
        type_c = history[history['value'] == 2]['is_right'].mean()
        accuracy = history['is_right'].mean()

        correct_score = history[history['is_right']==1]['value'].sum()
        total_score = history['value'].sum()
        score = correct_score / total_score
        new_row = pd.DataFrame({'is_right': [score],
                                "answer":str(correct_score) + '/' + str(total_score),
                                "question":type_a,
                                "choices":type_b,
                                "truth":type_c,
                                })
        # 将新行追加到DataFrame的末尾
        history = pd.concat([history, new_row], ignore_index=True)
        if prefix == 'GM_':
            appendix = '_eva_{}.xlsx'.format(self.each_rsl)
        else:
            appendix = '_script_{}_chat_{}.xlsx'.format(self.args.rsl,self.args.rchl)
        path_excel = os.path.join(path,prefix+self.script_manager.script_name+appendix)
        history.to_excel(path_excel, index=False,  encoding='utf-8-sig')
        print(path_excel)
        return path

    def evaluate(self):
        """
        evaluate the answer of the question
        TODA
        """
        self.evaluate_history = []
        for j in tqdm(range(len(self.agents))):
            data = self.script_manager.agents_final_result[self.agents[j].name]
            for i in tqdm(range(len(data))):
            # for i in range(len(data)):
                question = data['question'][i]
                if pd.isna(data['c'][i]):
                    choices = '\nA: '+ data['a'][i] + '\n' + 'B: '+ data['b'][i] + '\n'
                elif pd.isna(data['d'][i]):
                    choices = '\nA: '+ data['a'][i] + '\n' + 'B: '+ data['b'][i] + '\n' + 'C: '+ data['c'][i] + '\n'
                elif pd.isna(data['e'][i]):
                    choices = '\nA: '+ data['a'][i] + '\n' + 'B: '+ data['b'][i] + '\n' + 'C: '+ data['c'][i] + '\n' + 'D: '+ data['d'][i] + '\n'
                else:
                    choices = '\nA: '+ data['a'][i] + '\n' + 'B: '+ data['b'][i] + '\n' + 'C: '+ data['c'][i] + '\n' + 'D: '+ data['d'][i] + '\n' + 'E: '+ data['e'][i] + '\n'
                truth = data['truth'][i]
                current_script = self.agents[j].database.query(question, self.rsl)
                chat_history = self.chat_dataset.query(question, self.rchl)
                if data['type'][i] == 'a':
                    question_prompt = self.prompts['evaluate_ask_rap'].replace('{current_script}',current_script).replace('{chat_history}',chat_history).replace('{question}',question).replace('{choices}',choices)
                else:
                    question_prompt = self.prompts['evaluate_ask_rap_multi'].replace('{current_script}',current_script).replace('{chat_history}',chat_history).replace('{question}',question).replace('{choices}',choices)
                answer = self.agents[j].speak(question_prompt,True)
                is_right = 0
                format_answer = json_format(answer)
                try:
                    if data['type'][i] == 'b':
                        format_answer = format_answer['answer'].lower()
                        if truth.lower() in format_answer:
                            is_right = 1
                    else:
                        format_answer = format_answer['answer'].lower()
                        if format_answer in truth.lower():
                            is_right = 1
                except:
                    pass
                self.evaluate_history.append(EvaluateItem(tester=self.agents[j].name, question = question, choices=choices, answer=answer,truth = truth, is_right=is_right, value=data['value'][i]))

    def evaluate_gm(self, version=0):
        self.evaluate_history = []

        data = self.script_manager.agents_final_result['GM']
        for i in tqdm(range(len(data))):
        # for i in range(len(data)):
            question = data['question'][i]
            if pd.isna(data['c'][i]):
                choices = '\nA: '+ data['a'][i] + '\n' + 'B: '+ data['b'][i] + '\n'
            elif pd.isna(data['d'][i]):
                choices = '\nA: '+ data['a'][i] + '\n' + 'B: '+ data['b'][i] + '\n' + 'C: '+ data['c'][i] + '\n'
            elif pd.isna(data['e'][i]):
                choices = '\nA: '+ data['a'][i] + '\n' + 'B: '+ data['b'][i] + '\n' + 'C: '+ data['c'][i] + '\n' + 'D: '+ data['d'][i] + '\n'
            else:
                choices = '\nA: '+ data['a'][i] + '\n' + 'B: '+ data['b'][i] + '\n' + 'C: '+ data['c'][i] + '\n' + 'D: '+ data['d'][i] + '\n' + 'E: '+ data['e'][i] + '\n'
            truth = data['truth'][i]
            # current_script = self.GM_agent.database.query(question, self.rsl)
            if version == 0:
                self.version = 'average'
                relative_script = ''
                for j in range(len(self.agents)):
                    name = self.agents[j].name
                    current_script = self.agents[j].database.query(question, self.each_rsl)
                    relative_script +="The information in {}'s script related to this question is: \n {} \n\n".format(name, current_script)

            if data['type'][i] == 'a':
                question_prompt = self.prompts['evaluate_ask_rap_gm'].replace('{current_script}',relative_script).replace('{question}',question).replace('{choices}',choices)
            else:
                question_prompt = self.prompts['evaluate_ask_rap_multi_gm'].replace('{current_script}',relative_script).replace('{question}',question).replace('{choices}',choices)
            answer = self.GM_agent.speak(question_prompt, True)
            is_right = 0
            format_answer = json_format(answer)
            try:
                format_answer = format_answer['answer'].lower()
                if data['type'][i] == 'b':
                    if truth.lower() in format_answer:
                        is_right = 1
                else:
                    if format_answer[0] in truth.lower():
                        is_right = 1
            except:
                pass
            self.evaluate_history.append(EvaluateItem(tester="GM", question = question, choices=choices, answer=answer,truth = truth, is_right=is_right, value=data['value'][i]))


    def start(self):
        self.self_introduction_phase()
        # count = 0
        # while count < self.args.max_turn:
        #     self.open_talk()
        #     if self.args.constraint:
        #         self.constraints_space()
        #     count += 1
    def ablation_turn(self):
        self.self_introduction_phase()
        count = 0
        while count < self.args.max_turn:
            self.open_talk()

            self.evaluate()
            self.args.rsl = 5000
            self.args.rchl = 5000
            self.evaluate()
            self.save_evaluate_history(self.args.output_root_path, 'turn_'+str(count)+'_0')
            self.save_evaluate_history(self.args.output_root_path, 'turn_'+str(count)+'_1')
            self.save_evaluate_history(self.args.output_root_path, 'turn_'+str(count)+'_2')
            self.save_history(self.args.output_root_path, 'turn_'+str(count))
            # self.save_evaluate_history(self.args.output_root_path, 'turn_'+str(count)+'_3')
            # self.save_evaluate_history(self.args.output_root_path, 'turn_'+str(count)+'_4')

            if count + 1 == self.args.max_turn:
                continue
            if self.args.constraint:
                self.constraints_space()
            print(count)
            count += 1

    def save_args_to_json(self, args, filename):
        args_dict = vars(args)
        with open(filename, 'w') as file:
            json.dump(args_dict, file, indent=4, ensure_ascii=False)

    def mkdir_output_dir(self, args, post = ''):
        args.output_root_path = os.path.join(args.output_root_path, args.script_name+post)

        if not os.path.exists(args.output_root_path):
            os.makedirs(args.output_root_path)

    def save_params(self, args):
        input_token = []
        output_token = []
        input_token_eva = []
        output_token_eva = []

        for i in range(len(self.agents)):
            input_token += self.agents[i].session.input_token
            output_token += self.agents[i].session.output_token
            input_token_eva += self.agents[i].session.input_token_eva
            output_token_eva += self.agents[i].session.output_token_eva
        args.input_token = sum(input_token)
        args.output_token = sum(output_token)
        args.input_price = args.input_token * 0.0005 / 1000
        args.output_price = args.output_token * 0.0015 / 1000

        args.input_token_eva = sum(input_token_eva)
        args.output_token_eva = sum(output_token_eva)
        args.input_eva_price = args.input_token_eva * 0.0005 / 1000
        args.output_eva_price = args.output_token_eva * 0.0015 / 1000

        args.play_price = args.input_price + args.output_price
        args.eva_price = args.input_eva_price + args.output_eva_price

        self.save_args_to_json(args, os.path.join(args.output_root_path, 'args.json'))

    def mkdir_output_dir_time(self, args):
        now = datetime.now()
        time_string = now.strftime("%Y%m%d%H%M%S")
        args.output_root_path = os.path.join(args.output_root_path, args.script_name+time_string)
        self.args.output_root_path = args.output_root_path
        if not os.path.exists(args.output_root_path):
            os.makedirs(args.output_root_path)
        self.save_args_to_json(args, os.path.join(args.output_root_path, 'args.json'))

    def save_all_evaluate(self, path):

        path = os.path.join(path, 'evaluate')
        xlsx_files = glob.glob(os.path.join(path, "*.xlsx"))
        data = []
        for file in xlsx_files:
            df = pd.read_excel(file)
            last_value = df["is_right"].iloc[-1]
            detail_value = df["answer"].iloc[-1]
            type_a = df["question"].iloc[-1]
            type_b = df["choices"].iloc[-1]
            type_c = df["truth"].iloc[-1]

            filename = os.path.basename(file)
            data.append((filename,type_a, type_b,type_c,detail_value, last_value, len(df)-1))
        result_df = pd.DataFrame(data, columns=["filename", "type_a","type_b","type_c","detail_value", "score","length"])
        result_df = result_df.sort_values(by='filename')
        accuracy = result_df[~result_df['filename'].str.contains('GM') & ~result_df['filename'].str.contains('no_play')]['score'].mean()
        accuracy_a = result_df[~result_df['filename'].str.contains('GM') & ~result_df['filename'].str.contains('no_play')]['type_a'].mean()
        accuracy_b = result_df[~result_df['filename'].str.contains('GM') & ~result_df['filename'].str.contains('no_play')]['type_b'].mean()
        accuracy_c = result_df[~result_df['filename'].str.contains('GM') & ~result_df['filename'].str.contains('no_play')]['type_c'].mean()
        new_row = pd.DataFrame({'score': [accuracy],'type_a': [accuracy_a],'type_b': [accuracy_b],'type_c': [accuracy_c]})
        result_df = pd.concat([result_df, new_row], ignore_index=True)
        import numpy as np
        def format_significant(x):
            if isinstance(x, float):
                return np.format_float_positional(x, precision=3, unique=False, fractional=False, trim='k')
            return x

        result_df = result_df.applymap(format_significant)
        output_csv_path = os.path.join(path, "data_summary.csv")
        result_df.to_csv(output_csv_path, index=False, encoding='utf-8-sig')
        print(output_csv_path)

def evalute_para(game,path, prefix = ''):

    game.evaluate()
    game.save_evaluate_history(path, prefix)

def main():

    """play game"""
    scriptmanager = ScriptManager(args)
    game = Game(scriptmanager, args)
    game.mkdir_output_dir(args)
    game.start()

    game.save_history(args.output_root_path)

    path = args.output_root_path

    evalute_para(game, path, "0")
    evalute_para(game, path, "1")
    evalute_para(game, path, "2")

    game.save_all_evaluate(path)
    game.save_params(args)

def ablation():
    args.output_root_path = './log'

    """play game"""
    scriptmanager = ScriptManager(args)
    game = Game(scriptmanager, args)
    game.mkdir_output_dir_time(args)
    game.ablation_turn()

    game.save_all_evaluate(args.output_root_path)



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # for script
    parser.add_argument('--script_root', default=r"./chinese", type=str,required=False)
    parser.add_argument('--database_root', default=r"./database", type=str,required=False)
    parser.add_argument('--script_name', default="131-Sin (4 people closed)", type=str, required=False)
    parser.add_argument('--prompt_path', default="./prompts.yaml", type=str,required=False)
    parser.add_argument('--sensor_path', default="./sensors.json", type=str,required=False)

    parser.add_argument('--load_history', default="./", type=str,required=False)

    # for model
    parser.add_argument('--temperature', default=1, type=float, help='the diversity of generated text')
    parser.add_argument('--model_type', default="gpt35", type=str, required=False, help='[gpt35, gpt4, llama13b]')

    # for game

    parser.add_argument('--rsl', default=4000, type=float, help='the max length for relative script length (retrival from current script by RAG')
    parser.add_argument('--rchl', default=4000, type=float, help='the max length for relative chat history length (retrival from chat history by RAG')
    parser.add_argument('--each_rsl', default=3000, type=float, help='the max length for each agent\'s relative script length (retrival from current script by RAG')
    parser.add_argument('--question_number', default='1', type=str, help='the max length for each agent\'s relative script length (retrival from current script by RAG')
    parser.add_argument('--max_turn', default=3, type=int, help='the max length for each agent\'s relative script length (retrival from current script by RAG')
    parser.add_argument('--constraint', default=1, type=int, help='the max length for each agent\'s relative script length (retrival from current script by RAG')
    parser.add_argument('--sensor', default=1, type=int, help='the max length for each agent\'s relative script length (retrival from current script by RAG')
    parser.add_argument('--ablation', default=0, type=int, help='the max length for each agent\'s relative script length (retrival from current script by RAG')
    parser.add_argument('--evaluate', default=0, type=int, help='the max length for each agent\'s relative script length (retrival from current script by RAG')
    parser.add_argument('--is_english', default=1, type=int, help='the max length for each agent\'s relative script length (retrival from current script by RAG')

# TODA LLAMA
    parser.add_argument('--ckpt_dir', default="/scratch/prj/lmrep/llama2_model/Llama-2-13b-hf", type=str,required=False, help='if llama')

    # for log
    parser.add_argument('--output_root_path', default=r'./log_ours', type=str, required=False, help='max_tree_depth')

    args = parser.parse_args()
    agent_num2each_rsl={
        9: 1200,
        8: 1600,
        7: 2000,
        6: 2500,
        5: 3000,
        4: 3500,
    }

    model2ckpt = {
        "gemma7b": "path",
        "llama70b": "path",
        "llama13b": "path",
        "llama7b": "path",
        "gpt35": "",
        "gpt4": ""
    }
    if args.is_english:
        args.database_root = "./database_english"
        args.script_root = "./english"
    args.ckpt_dir = model2ckpt[args.model_type]
    if args.ablation:
        print('ablation begin',args.max_turn,args.question_number)
        ablation()
    elif args.evaluate:
        print('evalute begin')
        # evalute_from_history()
    else:
        main()