Simulator_LLM.py

import numpy as np
import pandas as pd

import openai
import requests
import time
import random

import re
from tqdm import tqdm
from itertools import combinations
import string

import sys,os

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../Functions')))
from Task import MAB_new
from DataProcessTools import DataProcessor

# retry decorator (https://cookbook.openai.com/examples/how_to_handle_rate_limits)
def retry_with_exponential_backoff(
    func,
    initial_delay: float = 1,
    exponential_base: float = 2,
    jitter: bool = True,
    max_retries: int = 10,
    errors: tuple = (openai.RateLimitError, KeyError),
):
    """Retry a function with exponential backoff."""

    def wrapper(*args, **kwargs):
        # Initialize variables
        num_retries = 0
        delay = initial_delay

        # Loop until a successful response or max_retries is hit or an exception is raised
        while True:
            try:
                return func(*args, **kwargs)

            # Retry on specified errors
            except errors as e:
                # Increment retries
                num_retries += 1

                # Check if max retries has been reached
                if num_retries > max_retries:
                    raise Exception(
                        f"Maximum number of retries ({max_retries}) exceeded."
                    )

                # Increment the delay
                delay *= exponential_base * (1 + jitter * random.random())

                # Sleep for the delay
                time.sleep(delay)

            # Raise exceptions for any errors not specified
            except Exception as e:
                raise e

    return wrapper

class simLLM_GCB_2AB:
    
    def __init__(self, name_model,setting_block, n_trials=50, client=None, hf_token=None, 
                 list_model_openAI = ["gpt-3.5-turbo", "gpt-4o-2024-11-20",'Qwen/Qwen2.5-72B-Instruct']):
        # set LLM
        self.name_model = name_model
        self.client = client
        self.hf_token = hf_token
        self.list_model_openAI = list_model_openAI

        self.n_trials = n_trials
        self.B_m0 = 0
        self.tolerance = 0.4

        self.mab_task = MAB_new(self.n_trials) # multi-armed bandit task      
        data_process = DataProcessor()

        self.name_options = list(string.ascii_uppercase) # list from which to choose 2 different letters for each condition

        self.setting_block = setting_block # 
        self.setting_blocklist_CorrectO1 = data_process.CorrectO1 # 

        # column name cn_ for dataframe
        self.cn_id_subj = data_process.ID_cn
        self.cn_condition = data_process.condition_cn
        self.cn_trial = data_process.trial_cn
        self.cn_feedback_O1 = data_process.O1_cn # feedback of O1
        self.cn_feedback_O2 = data_process.O2_cn # feedback of O2
        self.cn_ACC = data_process.ACC_cn # ACC definted by higher EV option
        self.cn_points = data_process.points_cn # points obtained for the trial


        self.cn_name_model = "name_LLM"
        self.cn_name_O1 = "name_O1"
        self.cn_name_O2 = "name_O2"
        self.cn_id_choice_LLM = "id_choice_LLM"
        self.cn_name_choice_LLM = "name_choice_LLM"
        self.cn_prompt = "prompt"
        self.cn_prompt_all_next = 'prompt_all_next'

        self.cn_name_est_LLM = 'name_est'
        self.cn_id_est_LLM = 'id_est'
        self.cn_ProbEst_LLM = 'prob_est_LLM'
        self.cn_AverEst_LLM = 'aver_est_LLM'
        self.cn_prompt_all_next = 'prompt_all_next'

        # self.kept_cols = [self.ID_cn,self.Session_cn, self.condition_cn, self.trial_cn, self.O1_cn, self.O2_cn,
        #             self.ACC_cn, self.RT_cn, self.points_cn, self.CT_cn, self.bonus_cn]

        self.text_first = f"In the first round, you choose slot machine:"
    
    def get_2letters(self,list_letters):
        return random.sample(list_letters, 2)  # get two different letters
    
    def sim_experiment(self,id_subj): # simulate all blocks, return behavioral dataframe
        self.id_subj = id_subj
        data_choices = pd.DataFrame()
        # data_estimations = pd.DataFrame()
        # random.shuffle(self.name_options)  

        # # check enough names for all block options
        # required_length = 2 * len(self.setting_block)
        # if len(self.name_options) < required_length:
        #     raise ValueError(f"self.name_options is too short! Current length: {len(self.name_options)}, required: {required_length}")

        for _, name_block in tqdm(enumerate(self.setting_block), desc="blocks"): # get block key word, and id_block for randomly select option name
            
            name_block = name_block
            name_O1O2 = self.get_2letters(self.name_options)
            name_O1 = name_O1O2[0]
            name_O2 = name_O1O2[1]
            # get B_m1s, B_p1s, G_ms, G_stds by block
            setting = self.setting_block[name_block]

            # get block simulation: list
            data_single_block = self.sim_block(name_block,name_O1,name_O2,**setting)
            data_single_block = pd.DataFrame(data_single_block)
            # data_single_block_estimation = pd.DataFrame(data_single_block_estimation)

            # concate data
            data_choices = pd.concat([data_choices, data_single_block], ignore_index=True)

            # remove estimation
            # data_estimations = pd.concat([data_estimations, data_single_block_estimation], ignore_index=True)

        return data_choices
    
    def sim_block(self,name_block,name_O1,name_O2,**kwargs): # get settings for each block, return resposne of one block/condition 

        B_m1s = kwargs.get('B_m1s', None)
        B_p1s = kwargs.get('B_p1s', None)
        G_ms = kwargs.get('G_ms', None)
        G_stds = kwargs.get('G_stds', None)

        # generate feedback list: [[O1_feedbacks],[O2_feedbacks]]
        feedbacks_O1O2 = self.mab_task.BimodalGaussian(B_m1s, B_p1s, G_ms, G_stds, B_m0 = self.B_m0,tolerance = self.tolerance)
        
        # multi-turn conversation
        data_single_block  = self.sim_trials(name_O1,name_O2,name_block,feedbacks_O1O2)

        return data_single_block
    
    def sim_trials(self,name_O1,name_O2,name_block,feedbacks_O1O2): # muti-turn conversation across trials, return response of each trial
        # simulation for a whole block
        data_trials = [] # decision making
        
        prompt = []

        # start block
        ## adding random option letters for init system
        text_content_system = self.generate_text_system(name_O1,name_O2)                      
        prompt.append(self.generate_prompt("system", text_content_system))

        ## adding prompt for the 1st trial
        prompt.append(self.generate_prompt("user",self.text_first))

        for t in tqdm(range(self.n_trials), desc="trials for " + name_block):

            id_trial = t+1 # current trial
            
            # make decision
            choice_LLM = self.query_model(prompt)[0] # occasionally, a '.' appears at the end

            # get feedback
            feedback_O1 = feedbacks_O1O2[0][t]
            feedback_O2 = feedbacks_O1O2[1][t]

            # generate trial history prompt
            ## conver choice to trial history prompt
            id_choice, name_chosen, feedback_chosen, name_unchosen, feedback_unchosen = self.convert_chosen_feedback(choice_LLM,name_O1,name_O2,feedback_O1,feedback_O2)
            ACC = self.convert_ACC(id_choice,name_block)
            text_next_trial = self.generate_text_trial(id_trial,name_chosen,feedback_chosen,name_unchosen,feedback_unchosen)
            ## add hisory to prompt
            prompt.append(self.generate_prompt("user",text_next_trial))


            # data
            data_singe_trial = {
                        self.cn_name_model: self.name_model,
                        self.cn_id_subj:self.id_subj,
                        self.cn_condition: name_block,
                        self.cn_trial: id_trial, # start from 1

                        self.cn_name_O1:name_O1,
                        self.cn_name_O2:name_O2,
                        self.cn_feedback_O1:feedback_O1,
                        self.cn_feedback_O2:feedback_O2,

                        self.cn_id_choice_LLM: id_choice,
                        self.cn_name_choice_LLM: name_chosen,
        
                        self.cn_ACC: ACC,
                        self.cn_points:feedback_chosen,
                        self.cn_prompt: text_next_trial,
                        self.cn_prompt_all_next: list(prompt)
            }
            ## 
            
            data_trials.append(data_singe_trial)
        return data_trials
    
    def sim_est_prob_aver(self, name_O1,name_O2,prompt):

        data_estimations = [] # probability or expected value estimation
                # probability and expected value estimation
        for id_choice_est, name_O in enumerate([name_O1, name_O2], start=1):  

            text_probability = self.generate_text_probability(name_O)
            prompt.append(self.generate_prompt("user",text_probability))
            prob_LLM = self.query_model(prompt)

            text_average = self.generate_text_average(name_O)
            prompt.append(self.generate_prompt("user",text_average))
            aver_LLM = self.query_model(prompt)

            data_single_estimation = {
                self.cn_name_est_LLM: name_O,
                self.cn_id_est_LLM: id_choice_est,
                self.cn_ProbEst_LLM: prob_LLM,
                self.cn_AverEst_LLM: aver_LLM
            }

            data_estimations.append(data_single_estimation)
        return data_estimations

    
    def generate_text_system(self, name_O1, name_O2):
        list_names = [name_O1, name_O2]
        random.shuffle(list_names)
        text_system = f"""
    You are playing a game with the goal of accumulating as much money as possible over the course of several rounds.
    There are two slot machines (slot machine {list_names[0]} and slot machine {list_names[1]}) that deliver money at different magnitudes and probabilities.
    One slot machine delivers more money than another on average.
    Your total payoff will be the cumulative sum of the money you win across all rounds of the game.
    Your goal is to maximize your total payoff by choosing between slot machines {list_names[0]} and {list_names[1]} strategically in each round.
    Please respond with either "{list_names[0]}" or "{list_names[1]}" based on the strategy that can maximize your total payoff. 
    Give your answer without explaining your reasoning.
"""
        return text_system
    
    def generate_text_trial(self,id_trial,name_chosen,feedback_chosen,name_unchosen,feedback_unchosen):
        text_trial = f"""
        In round {id_trial}, you chose slot machine {name_chosen} and got {feedback_chosen} dollars. The unchosen slot machine {name_unchosen} delivered {feedback_unchosen} dollars.
        In the next round, you choose slot machine:"""
    
        return text_trial    
    
    def generate_text_probability(self,name_O):
        text_probability = f"""The probability of getting any dollar from slot machine {name_O} is:"""
    
        return text_probability   
    
    def generate_text_average(self,name_O):
        text_average = f"""The average amount of dollars delivered by slot machine {name_O} is:"""
    
        return text_average 

    def generate_prompt(self,role,text_content):

        # check role
        valid_roles = ["system", "user", "assistant"]
        if role not in valid_roles:
            raise ValueError(f"Invalid role: {role}. Role must be one of {valid_roles}")

        message = {      
            "role": role,
            "content": [
                {
                "type": "text",
                "text": text_content        
                }
            ]
        }
        return message
    
    def convert_ACC(self,id_choice,name_block):
        if id_choice not in [1, 2, -1]:
            raise ValueError(f"Invalid id_choice: {id_choice}. Expected -1, 1 or 2.")

        if name_block in self.setting_blocklist_CorrectO1:
            return 1 if id_choice == 1 else 0
        else:
            return 0 if id_choice == 1 else 1
    
    def convert_chosen_feedback(self,choice_LLM,name_O1,name_O2,feedback_O1,feedback_O2):

        # conver choice_LLM to id_choice (O1 or O2); name; feedbacks
        if choice_LLM == name_O1:
            id_choice = 1
            name_chosen = name_O1
            feedback_chosen = feedback_O1

            name_unchosen = name_O2
            feedback_unchosen = feedback_O2

        elif choice_LLM == name_O2:
            id_choice = 2
            name_chosen = name_O2
            feedback_chosen = feedback_O2

            name_unchosen = name_O1
            feedback_unchosen = feedback_O1
        else: # if generate something else other than the only one letter, 
            id_choice = random.choice([1, 2])
            if id_choice == 1:
                name_chosen = name_O1
                feedback_chosen = feedback_O1
                name_unchosen = name_O2
                feedback_unchosen = feedback_O2
            else:
                name_chosen = name_O2
                feedback_chosen = feedback_O2
                name_unchosen = name_O1
                feedback_unchosen = feedback_O1
        
            id_choice = -1 #  return id_choice = -1 if randomly choose

        return id_choice, name_chosen, feedback_chosen, name_unchosen, feedback_unchosen

    @retry_with_exponential_backoff
    def query_model(self, messages, temperature=0, system=None):
        if self.name_model in self.list_model_openAI :
            response = self.client.chat.completions.create(
                model = self.name_model,
                messages = messages,
                response_format={"type": "text"},
                temperature = temperature)
            return response.choices[0].message.content.strip()
        ## TODO: change to model list
        # Hugging face API
        elif self.name_model.startswith('meta-llama') or self.name_model.startswith('mistralai'):
            API_URL = f"https://api-inference.huggingface.co/models/{self.name_model}"
            headers = {"Authorization": f"Bearer {self.hf_token}"}
            payload = {"inputs": messages, 
            "parameters": {"do_sample": True if temperature > 0 else False},
            "options": {"use_cache": False, "wait_for_model": True}}  #don't use cache- we want a fresh response
            response = requests.post(API_URL, headers=headers, json=payload)
            return response.json()[0]['generated_text'][len(messages):].strip()

class simLLM_EV_2AB_BG(simLLM_GCB_2AB):
    def generate_text_system(self, name_O1, name_O2):

        list_names = [name_O1, name_O2]
        random.shuffle(list_names)
        text_system = f"""
        You are playing a game with the goal of choosing the slot machine with higher total payout across all rounds of the game.
        There are two slot machines (slot machine {list_names[0]} and slot machine {list_names[1]}) that deliver money at different magnitudes and probabilities.
        One slot machine delivers higher total amount of money across all rounds at the end of the game. This is the correct slot machine to choose.
        Your total payoff will determined by the percentage of rounds in which you chose the correct option.
        Your goal is to maximize your accuracy based on the estimation of the average amount of money of each slot machine in each round.
        Please respond with either "{list_names[0]}" or "{list_names[1]}" based on the strategy that can maximize your total payoff. 
        Give your answer without explaining your reasoning. """
        return text_system

    def sim_block(self,name_block,name_O1,name_O2, RinT = 10,**kwargs): # get settings for each block, return resposne of one block/condition 

        '''
        Using RinT = 10, to ensure the expected value of the first 10 trials match the setting  
        '''

        B_m1s = kwargs.get('B_m1s', None)
        B_p1s = kwargs.get('B_p1s', None)
        G_ms = kwargs.get('G_ms', None)
        G_stds = kwargs.get('G_stds', None)

        # generate feedback list: [[O1_feedbacks],[O2_feedbacks]]
        feedbacks_O1O2 = self.mab_task.BimodalGaussian_RinT(B_m1s, B_p1s, G_ms, G_stds, B_m0 = self.B_m0,tolerance = self.tolerance, RinT = RinT)


        # multi-turn conversation
        data_single_block  = self.sim_trials(name_O1,name_O2,name_block,feedbacks_O1O2)

        return data_single_block