        obs = obs.split('\n')[0]
    obs = obs.lstrip()
    return obs

def instruction_parser(ins):
    ins = ins.encode().decode("unicode-escape").encode("latin1").decode("utf-8")
    ins = ins[13:].lstrip()
    return ins

def get_instruction(obs):
    obs = obs.split('\n')
    for idx, ob in enumerate(obs):
        if "Instruction:" in ob:
            return obs[idx+1].strip()

def get_buttons(obs):
  pattern = r"\[(.*?)\]"
  matches = re.findall(pattern, obs)
  return matches

def get_query(search):
  pattern = r"\[(.*?)\]"
  matches = re.findall(pattern, search)
  if matches:
    return matches[0]
  else:
    return search.replace('search', '')

def get_actions(obs):
  if "[Search]" in obs:
    avai_actions = {'search':[]}
  else:
    avai_actions = {'click':get_buttons(obs)}
  return avai_actions

def get_env_botton(bottons):
  if "Search" in bottons:
    avai_actions = {'search':[]}
  else:
    avai_actions = {'click':bottons}
  return avai_actions

def session_save(session_idx, actions, observations, retrieved_items, file_path):
  save_dict_list = []
  for idx in range(len(actions)):
    save_dict =  {"step_id": f"session_{session_idx}_step_{idx}", "action": actions[idx], "observation":observations[idx], "retrieved_items": retrieved_items[idx]}
    save_dict_list.append(save_dict)


def get_instruction(obs):
    obs = obs.split('\n')
    for idx, ob in enumerate(obs):
        if "Instruction:" in ob:
            return obs[idx+1].strip()

def get_buttons(obs):
  pattern = r"\[(.*?)\]"
  matches = re.findall(pattern, obs)
  return matches

def get_query(search):
  pattern = r"\[(.*?)\]"
  matches = re.findall(pattern, search)
  if matches:
    return matches[0]
  else:
    return search.replace('search', '')

def get_actions(obs):
  if "[Search]" in obs:
    avai_actions = {'search':[]}
  else:
    avai_actions = {'click':get_buttons(obs)}
  return avai_actions

def get_env_botton(bottons):
  if "Search" in bottons:
    avai_actions = {'search':[]}
  else:
    avai_actions = {'click':bottons}
  return avai_actions

def session_save(session_idx, actions, observations, retrieved_items, file_path):
  save_dict_list = []
  for idx in range(len(actions)):
    save_dict =  {"step_id": f"session_{session_idx}_step_{idx}", "action": actions[idx], "observation":observations[idx], "retrieved_items": retrieved_items[idx]}
    save_dict_list.append(save_dict)
  # get goal
  goal_file_name = f"./webshop/user_session_logs/mturk/{session_idx}.jsonl"
  with open(goal_file_name) as f:
    for line in f:
      page = json.loads(line)
      break
    goal = page["goal"]


def get_buttons(obs):
  pattern = r"\[(.*?)\]"
  matches = re.findall(pattern, obs)
  return matches

def get_query(search):
  pattern = r"\[(.*?)\]"
  matches = re.findall(pattern, search)
  if matches:
    return matches[0]
  else:
    return search.replace('search', '')

def get_actions(obs):
  if "[Search]" in obs:
    avai_actions = {'search':[]}
  else:
    avai_actions = {'click':get_buttons(obs)}
  return avai_actions

def get_env_botton(bottons):
  if "Search" in bottons:
    avai_actions = {'search':[]}
  else:
    avai_actions = {'click':bottons}
  return avai_actions

def session_save(session_idx, actions, observations, retrieved_items, file_path):
  save_dict_list = []
  for idx in range(len(actions)):
    save_dict =  {"step_id": f"session_{session_idx}_step_{idx}", "action": actions[idx], "observation":observations[idx], "retrieved_items": retrieved_items[idx]}
    save_dict_list.append(save_dict)
  # get goal
  goal_file_name = f"./webshop/user_session_logs/mturk/{session_idx}.jsonl"
  with open(goal_file_name) as f:
    for line in f:
      page = json.loads(line)
      break
    goal = page["goal"]
  # saving
  save_dict_list = [goal] + save_dict_list
  with open(file_path, 'a') as f:
    json.dump(save_dict_list, f)
    f.write('\n')


  matches = re.findall(pattern, search)
  if matches:
    return matches[0]
  else:
    return search.replace('search', '')

def get_actions(obs):
  if "[Search]" in obs:
    avai_actions = {'search':[]}
  else:
    avai_actions = {'click':get_buttons(obs)}
  return avai_actions

def get_env_botton(bottons):
  if "Search" in bottons:
    avai_actions = {'search':[]}
  else:
    avai_actions = {'click':bottons}
  return avai_actions

def session_save(session_idx, actions, observations, retrieved_items, file_path):
  save_dict_list = []
  for idx in range(len(actions)):
    save_dict =  {"step_id": f"session_{session_idx}_step_{idx}", "action": actions[idx], "observation":observations[idx], "retrieved_items": retrieved_items[idx]}
    save_dict_list.append(save_dict)
  # get goal
  goal_file_name = f"./webshop/user_session_logs/mturk/{session_idx}.jsonl"
  with open(goal_file_name) as f:
    for line in f:
      page = json.loads(line)
      break
    goal = page["goal"]
  # saving
  save_dict_list = [goal] + save_dict_list
  with open(file_path, 'a') as f:
    json.dump(save_dict_list, f)
    f.write('\n')


"""
 Copyright (c) 2023, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: Apache License 2.0
 For full license text, see the LICENSE file in the repo root or https://www.apache.org/licenses/LICENSE-2.0
"""

import time
import os
import web_run.pre_prompt as pre_prompt
from web_run.utils import get_query
from web_run.llms import token_enc

control_dict = {"Search_Webrun_Agent":"search",
                "Think_Webrun_Agent":"think",
                "Click_Webrun_Agent":"click"}

class ClickAgent:
    def __init__(self, llm, context_len=2000):
        self.type = "Click_Webrun_Agent"
        self.life_label = time.time()
        self.name = f"{self.type}_{self.life_label}"
        self.llm = llm
        self.context_len = context_len
        self.task = None

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        if nor_text.startswith('click'):
            if 'click' in available_actions:
                for a in available_actions['click']:
                    if a.lower() in nor_text:
                        return f"click[{a}]"
        return text

    def avai_action_prompt(self, available_actions):
        if 'search' in available_actions:
            return "{search}"
        elif 'click' in available_actions:
            actions = available_actions['click']
            ac_string = "\t".join(actions)
            return f"click:[{ac_string}]"

    def prompt_layer(self, control_prompt, available_actions):
        one_shot = pre_prompt.click
        initial = f"Instruction:\n{self.task}"
        actions_prompt = f"current available action is {self.avai_action_prompt(available_actions)[:3000]}"
        remain_context_space = (self.context_len - len(token_enc.encode(initial+one_shot+actions_prompt)))*3
        prompt = f"{one_shot}\n{initial}\n{control_prompt[-int(remain_context_space):]}\n{actions_prompt}\n\nAction:"
        if len(token_enc.encode(prompt)) > self.context_len:
            print("this prompt may be over context length requirements:")
            print("---------\n")
            print(prompt)
            print("\n---------")
        return prompt

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, control_prompt, available_actions):
        prompt = self.prompt_layer(control_prompt, available_actions)
        action = self.llm_layer(prompt).lstrip(' ')
        action = self.action_parser(action, available_actions)
        return action

class SearchAgent:
    def __init__(self, llm, context_len=2000):
        self.type = "Search_Webrun_Agent"
        self.life_label = time.time()
        self.name = f"{self.type}_{self.life_label}"
        self.llm = llm
        self.context_len = context_len
        self.task = None

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        return text

    def prompt_layer(self, control_prompt, available_actions):
        one_shot = pre_prompt.search
        remain_context_space = (self.context_len - len(token_enc.encode(one_shot)))*3
        prompt = f"{one_shot}{control_prompt[-int(remain_context_space):]}\n\nAction:"
        if len(token_enc.encode(prompt)) > self.context_len:
            print("this prompt may be over context length requirements:")
            print("---------\n")
            print(prompt)
            print("\n---------")
        return prompt

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, control_prompt, available_actions):
        prompt = self.prompt_layer(control_prompt, available_actions)
        action = self.llm_layer(prompt).lstrip(' ')
        action = self.action_parser(action, available_actions)
        return action

class ThinkAgent:
    def __init__(self, llm, context_len=2000):
        self.type = "Think_Webrun_Agent"
        self.name = f"{self.type}_{self.life_label}"
        self.llm = llm
        self.context_len = context_len
        self.task = None

    def action_parser(self, text):
        nor_text = text.strip().lower()
        if nor_text.startswith('think'):
            query = get_query(nor_text)
            return f"think[{query}]"
        return text

    def prompt_layer(self, agent, available_actions):
        one_shot = pre_prompt.click
        prompt = f"{one_shot}{self.observations[self.cur_session][0]}\n\nAction:"
        actions_prompt = f"current available action is {self.avai_action_prompt(available_actions)}"
        prompt = f"\n\nAction:"
        return prompt

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, control_prompt, available_actions):
        prompt = self.prompt_layer(control_prompt, available_actions)
        action = self.llm_layer(prompt).lstrip(' ')
        action = self.action_parser(action, available_actions)
        return action

class ControlAgent:
    def __init__(self, agents):
        self.type = "Search_Click_Control_Webrun_Agent"
        self.agents = agents
        self.task = {}
        self.actions = {}
        self.observations = {}
        self.agent_calls = {} # saving the agents types in sessions
        self.item_recall = {}
        self.cur_session = 0
        self.agents = agents
        self.control_dict = control_dict

    def new_session(self, idx, task):
        self.cur_session = idx
        self.task[idx] = task
        self.actions[idx] = ['reset']
        self.observations[idx] = []
        self.item_recall[idx] = []
        self.agent_calls[idx] = ["Search_Click_Controller"]
        self.task_assign(task)

    def task_assign(self, task):
        for agent in self.agents:
            agent.task = task

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        if nor_text.startswith('click'):
            if 'click' in available_actions:
                for a in available_actions['click']:
                    if a.lower() in nor_text:
                        return f"click[{a}]"
        return text

    def get_agents_types(self):
        return [a.type for a in self.agents]

    def prompt_layer(self, agent, available_actions):
        if agent.type == "Search_Webrun_Agent":
            return self.ask_search_agent()
        if agent.type == "Click_Webrun_Agent":
            return self.ask_click_agent()
        return 0

    def ask_search_agent(self):
        initial = self.task[self.cur_session]
        prompt = f"Instruction:\n{initial}\n"
        search_label = False
        for idx, agent_type in enumerate(self.agent_calls[self.cur_session]):
            if agent_type == "Search_Webrun_Agent":
                search_label = True
                action = self.actions[self.cur_session][idx]
                prompt += f"You have already {action}.\n"
        if search_label:
            prompt += "Can you change your search query?"
        return prompt

    def ask_click_agent(self):
        prompt = ""
        for idx in range(1, len(self.actions[self.cur_session])):
            agent_type = self.agent_calls[self.cur_session][idx]
            action = self.actions[self.cur_session][idx]
            observation = self.observations[self.cur_session][idx]
            if agent_type == "Click_Webrun_Agent":
                prompt += f"\n\nAction:{action}\nObservation:{observation}"
            else:
                prompt += f"\n\n{agent_type}:{action}\nObservation:{observation}"
        return prompt

    def call_agent(self, observation, available_actions):
        for agent in self.agents:
            if self.control_dict[agent.type] in available_actions:
                return agent

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, observation, available_actions=None):
        self.observations[self.cur_session].append(observation)
        agent = self.call_agent(observation, available_actions)
        prompt = self.prompt_layer(agent, available_actions)
        action = agent.forward(prompt, available_actions).lstrip()
        action = self.action_parser(action, available_actions)
        self.actions[self.cur_session].append(action)
        self.agent_calls[self.cur_session].append(agent.type)
        return action


ClickAgent

"""
 Copyright (c) 2023, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: Apache License 2.0
 For full license text, see the LICENSE file in the repo root or https://www.apache.org/licenses/LICENSE-2.0
"""

import time
import os
import web_run.pre_prompt as pre_prompt
from web_run.utils import get_query
from web_run.llms import token_enc

control_dict = {"Search_Webrun_Agent":"search",
                "Think_Webrun_Agent":"think",
                "Click_Webrun_Agent":"click"}

class ClickAgent:
    def __init__(self, llm, context_len=2000):
        self.type = "Click_Webrun_Agent"
        self.life_label = time.time()
        self.name = f"{self.type}_{self.life_label}"
        self.llm = llm
        self.context_len = context_len
        self.task = None

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        if nor_text.startswith('click'):
            if 'click' in available_actions:
                for a in available_actions['click']:
                    if a.lower() in nor_text:
                        return f"click[{a}]"
        return text

    def avai_action_prompt(self, available_actions):
        if 'search' in available_actions:
            return "{search}"
        elif 'click' in available_actions:
            actions = available_actions['click']
            ac_string = "\t".join(actions)
            return f"click:[{ac_string}]"

    def prompt_layer(self, control_prompt, available_actions):
        one_shot = pre_prompt.click
        initial = f"Instruction:\n{self.task}"
        actions_prompt = f"current available action is {self.avai_action_prompt(available_actions)[:3000]}"
        remain_context_space = (self.context_len - len(token_enc.encode(initial+one_shot+actions_prompt)))*3
        prompt = f"{one_shot}\n{initial}\n{control_prompt[-int(remain_context_space):]}\n{actions_prompt}\n\nAction:"
        if len(token_enc.encode(prompt)) > self.context_len:
            print("this prompt may be over context length requirements:")
            print("---------\n")
            print(prompt)
            print("\n---------")
        return prompt

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, control_prompt, available_actions):
        prompt = self.prompt_layer(control_prompt, available_actions)
        action = self.llm_layer(prompt).lstrip(' ')
        action = self.action_parser(action, available_actions)
        return action

class SearchAgent:
    def __init__(self, llm, context_len=2000):
        self.type = "Search_Webrun_Agent"
        self.life_label = time.time()
        self.name = f"{self.type}_{self.life_label}"
        self.llm = llm
        self.context_len = context_len
        self.task = None

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        return text

    def prompt_layer(self, control_prompt, available_actions):
        one_shot = pre_prompt.search
        remain_context_space = (self.context_len - len(token_enc.encode(one_shot)))*3


SearchAgent

        initial = f"Instruction:\n{self.task}"
        actions_prompt = f"current available action is {self.avai_action_prompt(available_actions)[:3000]}"
        remain_context_space = (self.context_len - len(token_enc.encode(initial+one_shot+actions_prompt)))*3
        prompt = f"{one_shot}\n{initial}\n{control_prompt[-int(remain_context_space):]}\n{actions_prompt}\n\nAction:"
        if len(token_enc.encode(prompt)) > self.context_len:
            print("this prompt may be over context length requirements:")
            print("---------\n")
            print(prompt)
            print("\n---------")
        return prompt

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, control_prompt, available_actions):
        prompt = self.prompt_layer(control_prompt, available_actions)
        action = self.llm_layer(prompt).lstrip(' ')
        action = self.action_parser(action, available_actions)
        return action

class SearchAgent:
    def __init__(self, llm, context_len=2000):
        self.type = "Search_Webrun_Agent"
        self.life_label = time.time()
        self.name = f"{self.type}_{self.life_label}"
        self.llm = llm
        self.context_len = context_len
        self.task = None

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        return text

    def prompt_layer(self, control_prompt, available_actions):
        one_shot = pre_prompt.search
        remain_context_space = (self.context_len - len(token_enc.encode(one_shot)))*3
        prompt = f"{one_shot}{control_prompt[-int(remain_context_space):]}\n\nAction:"
        if len(token_enc.encode(prompt)) > self.context_len:
            print("this prompt may be over context length requirements:")
            print("---------\n")
            print(prompt)
            print("\n---------")
        return prompt

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, control_prompt, available_actions):
        prompt = self.prompt_layer(control_prompt, available_actions)
        action = self.llm_layer(prompt).lstrip(' ')
        action = self.action_parser(action, available_actions)
        return action

class ThinkAgent:
    def __init__(self, llm, context_len=2000):
        self.type = "Think_Webrun_Agent"
        self.name = f"{self.type}_{self.life_label}"
        self.llm = llm
        self.context_len = context_len
        self.task = None

    def action_parser(self, text):
        nor_text = text.strip().lower()
        if nor_text.startswith('think'):
            query = get_query(nor_text)
            return f"think[{query}]"
        return text

    def prompt_layer(self, agent, available_actions):
        one_shot = pre_prompt.click
        prompt = f"{one_shot}{self.observations[self.cur_session][0]}\n\nAction:"
        actions_prompt = f"current available action is {self.avai_action_prompt(available_actions)}"


ThinkAgent

    def prompt_layer(self, control_prompt, available_actions):
        one_shot = pre_prompt.search
        remain_context_space = (self.context_len - len(token_enc.encode(one_shot)))*3
        prompt = f"{one_shot}{control_prompt[-int(remain_context_space):]}\n\nAction:"
        if len(token_enc.encode(prompt)) > self.context_len:
            print("this prompt may be over context length requirements:")
            print("---------\n")
            print(prompt)
            print("\n---------")
        return prompt

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, control_prompt, available_actions):
        prompt = self.prompt_layer(control_prompt, available_actions)
        action = self.llm_layer(prompt).lstrip(' ')
        action = self.action_parser(action, available_actions)
        return action

class ThinkAgent:
    def __init__(self, llm, context_len=2000):
        self.type = "Think_Webrun_Agent"
        self.name = f"{self.type}_{self.life_label}"
        self.llm = llm
        self.context_len = context_len
        self.task = None

    def action_parser(self, text):
        nor_text = text.strip().lower()
        if nor_text.startswith('think'):
            query = get_query(nor_text)
            return f"think[{query}]"
        return text

    def prompt_layer(self, agent, available_actions):
        one_shot = pre_prompt.click
        prompt = f"{one_shot}{self.observations[self.cur_session][0]}\n\nAction:"
        actions_prompt = f"current available action is {self.avai_action_prompt(available_actions)}"
        prompt = f"\n\nAction:"
        return prompt

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, control_prompt, available_actions):
        prompt = self.prompt_layer(control_prompt, available_actions)
        action = self.llm_layer(prompt).lstrip(' ')
        action = self.action_parser(action, available_actions)
        return action

class ControlAgent:
    def __init__(self, agents):
        self.type = "Search_Click_Control_Webrun_Agent"
        self.agents = agents
        self.task = {}
        self.actions = {}
        self.observations = {}
        self.agent_calls = {} # saving the agents types in sessions
        self.item_recall = {}
        self.cur_session = 0
        self.agents = agents
        self.control_dict = control_dict

    def new_session(self, idx, task):
        self.cur_session = idx
        self.task[idx] = task
        self.actions[idx] = ['reset']
        self.observations[idx] = []
        self.item_recall[idx] = []


ControlAgent

            return f"think[{query}]"
        return text

    def prompt_layer(self, agent, available_actions):
        one_shot = pre_prompt.click
        prompt = f"{one_shot}{self.observations[self.cur_session][0]}\n\nAction:"
        actions_prompt = f"current available action is {self.avai_action_prompt(available_actions)}"
        prompt = f"\n\nAction:"
        return prompt

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, control_prompt, available_actions):
        prompt = self.prompt_layer(control_prompt, available_actions)
        action = self.llm_layer(prompt).lstrip(' ')
        action = self.action_parser(action, available_actions)
        return action

class ControlAgent:
    def __init__(self, agents):
        self.type = "Search_Click_Control_Webrun_Agent"
        self.agents = agents
        self.task = {}
        self.actions = {}
        self.observations = {}
        self.agent_calls = {} # saving the agents types in sessions
        self.item_recall = {}
        self.cur_session = 0
        self.agents = agents
        self.control_dict = control_dict

    def new_session(self, idx, task):
        self.cur_session = idx
        self.task[idx] = task
        self.actions[idx] = ['reset']
        self.observations[idx] = []
        self.item_recall[idx] = []
        self.agent_calls[idx] = ["Search_Click_Controller"]
        self.task_assign(task)

    def task_assign(self, task):
        for agent in self.agents:
            agent.task = task

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        if nor_text.startswith('click'):
            if 'click' in available_actions:
                for a in available_actions['click']:
                    if a.lower() in nor_text:
                        return f"click[{a}]"
        return text

    def get_agents_types(self):
        return [a.type for a in self.agents]

    def prompt_layer(self, agent, available_actions):
        if agent.type == "Search_Webrun_Agent":
            return self.ask_search_agent()
        if agent.type == "Click_Webrun_Agent":
            return self.ask_click_agent()
        return 0

    def ask_search_agent(self):
        initial = self.task[self.cur_session]
        prompt = f"Instruction:\n{initial}\n"
        search_label = False
        for idx, agent_type in enumerate(self.agent_calls[self.cur_session]):
            if agent_type == "Search_Webrun_Agent":
                search_label = True
                action = self.actions[self.cur_session][idx]
                prompt += f"You have already {action}.\n"
        if search_label:
            prompt += "Can you change your search query?"
        return prompt

    def ask_click_agent(self):
        prompt = ""
        for idx in range(1, len(self.actions[self.cur_session])):
            agent_type = self.agent_calls[self.cur_session][idx]
            action = self.actions[self.cur_session][idx]
            observation = self.observations[self.cur_session][idx]
            if agent_type == "Click_Webrun_Agent":
                prompt += f"\n\nAction:{action}\nObservation:{observation}"
            else:
                prompt += f"\n\n{agent_type}:{action}\nObservation:{observation}"
        return prompt

    def call_agent(self, observation, available_actions):
        for agent in self.agents:
            if self.control_dict[agent.type] in available_actions:
                return agent

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, observation, available_actions=None):
        self.observations[self.cur_session].append(observation)
        agent = self.call_agent(observation, available_actions)
        prompt = self.prompt_layer(agent, available_actions)
        action = agent.forward(prompt, available_actions).lstrip()
        action = self.action_parser(action, available_actions)
        self.actions[self.cur_session].append(action)
        self.agent_calls[self.cur_session].append(agent.type)
        return action


        return text

    def prompt_layer(self, agent, available_actions):
        one_shot = pre_prompt.click
        prompt = f"{one_shot}{self.observations[self.cur_session][0]}\n\nAction:"
        actions_prompt = f"current available action is {self.avai_action_prompt(available_actions)}"
        prompt = f"\n\nAction:"
        return prompt

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, control_prompt, available_actions):
        prompt = self.prompt_layer(control_prompt, available_actions)
        action = self.llm_layer(prompt).lstrip(' ')
        action = self.action_parser(action, available_actions)
        return action

class ControlAgent:
    def __init__(self, agents):
        self.type = "Search_Click_Control_Webrun_Agent"
        self.agents = agents
        self.task = {}
        self.actions = {}
        self.observations = {}
        self.agent_calls = {} # saving the agents types in sessions
        self.item_recall = {}
        self.cur_session = 0
        self.agents = agents
        self.control_dict = control_dict

    def new_session(self, idx, task):
        self.cur_session = idx
        self.task[idx] = task
        self.actions[idx] = ['reset']
        self.observations[idx] = []
        self.item_recall[idx] = []
        self.agent_calls[idx] = ["Search_Click_Controller"]
        self.task_assign(task)

    def task_assign(self, task):
        for agent in self.agents:
            agent.task = task

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        if nor_text.startswith('click'):


        self.observations = {}
        self.agent_calls = {} # saving the agents types in sessions
        self.item_recall = {}
        self.cur_session = 0
        self.agents = agents
        self.control_dict = control_dict

    def new_session(self, idx, task):
        self.cur_session = idx
        self.task[idx] = task
        self.actions[idx] = ['reset']
        self.observations[idx] = []
        self.item_recall[idx] = []
        self.agent_calls[idx] = ["Search_Click_Controller"]
        self.task_assign(task)

    def task_assign(self, task):
        for agent in self.agents:
            agent.task = task

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        if nor_text.startswith('click'):
            if 'click' in available_actions:
                for a in available_actions['click']:
                    if a.lower() in nor_text:
                        return f"click[{a}]"
        return text

    def get_agents_types(self):
        return [a.type for a in self.agents]

    def prompt_layer(self, agent, available_actions):
        if agent.type == "Search_Webrun_Agent":
            return self.ask_search_agent()
        if agent.type == "Click_Webrun_Agent":
            return self.ask_click_agent()
        return 0

    def ask_search_agent(self):
        initial = self.task[self.cur_session]
        prompt = f"Instruction:\n{initial}\n"
        search_label = False
        for idx, agent_type in enumerate(self.agent_calls[self.cur_session]):
            if agent_type == "Search_Webrun_Agent":
                search_label = True
                action = self.actions[self.cur_session][idx]
                prompt += f"You have already {action}.\n"


    def __init__(self, llm, context_len=2000):
        self.type = "Click_Webrun_Agent"
        self.life_label = time.time()
        self.name = f"{self.type}_{self.life_label}"
        self.llm = llm
        self.context_len = context_len
        self.task = None

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        if nor_text.startswith('click'):
            if 'click' in available_actions:
                for a in available_actions['click']:
                    if a.lower() in nor_text:
                        return f"click[{a}]"
        return text

    def avai_action_prompt(self, available_actions):
        if 'search' in available_actions:
            return "{search}"
        elif 'click' in available_actions:
            actions = available_actions['click']
            ac_string = "\t".join(actions)
            return f"click:[{ac_string}]"

    def prompt_layer(self, control_prompt, available_actions):
        one_shot = pre_prompt.click
        initial = f"Instruction:\n{self.task}"
        actions_prompt = f"current available action is {self.avai_action_prompt(available_actions)[:3000]}"
        remain_context_space = (self.context_len - len(token_enc.encode(initial+one_shot+actions_prompt)))*3
        prompt = f"{one_shot}\n{initial}\n{control_prompt[-int(remain_context_space):]}\n{actions_prompt}\n\nAction:"
        if len(token_enc.encode(prompt)) > self.context_len:
            print("this prompt may be over context length requirements:")
            print("---------\n")
            print(prompt)
            print("\n---------")
        return prompt

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, control_prompt, available_actions):
        prompt = self.prompt_layer(control_prompt, available_actions)
        action = self.llm_layer(prompt).lstrip(' ')


    def task_assign(self, task):
        for agent in self.agents:
            agent.task = task

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        if nor_text.startswith('click'):
            if 'click' in available_actions:
                for a in available_actions['click']:
                    if a.lower() in nor_text:
                        return f"click[{a}]"
        return text

    def get_agents_types(self):
        return [a.type for a in self.agents]

    def prompt_layer(self, agent, available_actions):
        if agent.type == "Search_Webrun_Agent":
            return self.ask_search_agent()
        if agent.type == "Click_Webrun_Agent":
            return self.ask_click_agent()
        return 0

    def ask_search_agent(self):
        initial = self.task[self.cur_session]
        prompt = f"Instruction:\n{initial}\n"
        search_label = False
        for idx, agent_type in enumerate(self.agent_calls[self.cur_session]):
            if agent_type == "Search_Webrun_Agent":
                search_label = True
                action = self.actions[self.cur_session][idx]
                prompt += f"You have already {action}.\n"
        if search_label:
            prompt += "Can you change your search query?"
        return prompt

    def ask_click_agent(self):
        prompt = ""
        for idx in range(1, len(self.actions[self.cur_session])):
            agent_type = self.agent_calls[self.cur_session][idx]
            action = self.actions[self.cur_session][idx]
            observation = self.observations[self.cur_session][idx]


            prompt += "Can you change your search query?"
        return prompt

    def ask_click_agent(self):
        prompt = ""
        for idx in range(1, len(self.actions[self.cur_session])):
            agent_type = self.agent_calls[self.cur_session][idx]
            action = self.actions[self.cur_session][idx]
            observation = self.observations[self.cur_session][idx]
            if agent_type == "Click_Webrun_Agent":
                prompt += f"\n\nAction:{action}\nObservation:{observation}"
            else:
                prompt += f"\n\n{agent_type}:{action}\nObservation:{observation}"
        return prompt

    def call_agent(self, observation, available_actions):
        for agent in self.agents:
            if self.control_dict[agent.type] in available_actions:
                return agent

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, observation, available_actions=None):
        self.observations[self.cur_session].append(observation)
        agent = self.call_agent(observation, available_actions)
        prompt = self.prompt_layer(agent, available_actions)
        action = agent.forward(prompt, available_actions).lstrip()
        action = self.action_parser(action, available_actions)
        self.actions[self.cur_session].append(action)
        self.agent_calls[self.cur_session].append(agent.type)
        return action


forward

    def ask_click_agent(self):
        prompt = ""
        for idx in range(1, len(self.actions[self.cur_session])):
            agent_type = self.agent_calls[self.cur_session][idx]
            action = self.actions[self.cur_session][idx]
            observation = self.observations[self.cur_session][idx]
            if agent_type == "Click_Webrun_Agent":
                prompt += f"\n\nAction:{action}\nObservation:{observation}"
            else:
                prompt += f"\n\n{agent_type}:{action}\nObservation:{observation}"
        return prompt

    def call_agent(self, observation, available_actions):
        for agent in self.agents:
            if self.control_dict[agent.type] in available_actions:
                return agent

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, observation, available_actions=None):
        self.observations[self.cur_session].append(observation)
        agent = self.call_agent(observation, available_actions)
        prompt = self.prompt_layer(agent, available_actions)
        action = agent.forward(prompt, available_actions).lstrip()
        action = self.action_parser(action, available_actions)
        self.actions[self.cur_session].append(action)
        self.agent_calls[self.cur_session].append(agent.type)
        return action


    def forward(self, control_prompt, available_actions):
        prompt = self.prompt_layer(control_prompt, available_actions)
        action = self.llm_layer(prompt).lstrip(' ')
        action = self.action_parser(action, available_actions)
        return action

class ControlAgent:
    def __init__(self, agents):
        self.type = "Search_Click_Control_Webrun_Agent"
        self.agents = agents
        self.task = {}
        self.actions = {}
        self.observations = {}
        self.agent_calls = {} # saving the agents types in sessions
        self.item_recall = {}
        self.cur_session = 0
        self.agents = agents
        self.control_dict = control_dict

    def new_session(self, idx, task):
        self.cur_session = idx
        self.task[idx] = task
        self.actions[idx] = ['reset']
        self.observations[idx] = []
        self.item_recall[idx] = []
        self.agent_calls[idx] = ["Search_Click_Controller"]
        self.task_assign(task)

    def task_assign(self, task):
        for agent in self.agents:
            agent.task = task

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        if nor_text.startswith('click'):
            if 'click' in available_actions:
                for a in available_actions['click']:
                    if a.lower() in nor_text:
                        return f"click[{a}]"
        return text

    def get_agents_types(self):
        return [a.type for a in self.agents]


        self.type = "Search_Click_Control_Webrun_Agent"
        self.agents = agents
        self.task = {}
        self.actions = {}
        self.observations = {}
        self.agent_calls = {} # saving the agents types in sessions
        self.item_recall = {}
        self.cur_session = 0
        self.agents = agents
        self.control_dict = control_dict

    def new_session(self, idx, task):
        self.cur_session = idx
        self.task[idx] = task
        self.actions[idx] = ['reset']
        self.observations[idx] = []
        self.item_recall[idx] = []
        self.agent_calls[idx] = ["Search_Click_Controller"]
        self.task_assign(task)

    def task_assign(self, task):
        for agent in self.agents:
            agent.task = task

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        if nor_text.startswith('click'):
            if 'click' in available_actions:
                for a in available_actions['click']:
                    if a.lower() in nor_text:
                        return f"click[{a}]"
        return text

    def get_agents_types(self):
        return [a.type for a in self.agents]

    def prompt_layer(self, agent, available_actions):
        if agent.type == "Search_Webrun_Agent":
            return self.ask_search_agent()
        if agent.type == "Click_Webrun_Agent":


        self.item_recall[idx] = []
        self.agent_calls[idx] = ["Search_Click_Controller"]
        self.task_assign(task)

    def task_assign(self, task):
        for agent in self.agents:
            agent.task = task

    def action_parser(self, text, available_actions):
        nor_text = text.strip().lower()
        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        if nor_text.startswith('click'):
            if 'click' in available_actions:
                for a in available_actions['click']:
                    if a.lower() in nor_text:
                        return f"click[{a}]"
        return text

    def get_agents_types(self):
        return [a.type for a in self.agents]

    def prompt_layer(self, agent, available_actions):
        if agent.type == "Search_Webrun_Agent":
            return self.ask_search_agent()
        if agent.type == "Click_Webrun_Agent":
            return self.ask_click_agent()
        return 0

    def ask_search_agent(self):
        initial = self.task[self.cur_session]
        prompt = f"Instruction:\n{initial}\n"
        search_label = False
        for idx, agent_type in enumerate(self.agent_calls[self.cur_session]):
            if agent_type == "Search_Webrun_Agent":
                search_label = True
                action = self.actions[self.cur_session][idx]
                prompt += f"You have already {action}.\n"
        if search_label:
            prompt += "Can you change your search query?"
        return prompt


        if nor_text.startswith('search'):
            query = get_query(nor_text)
            return f"search[{query}]"
        if nor_text.startswith('click'):
            if 'click' in available_actions:
                for a in available_actions['click']:
                    if a.lower() in nor_text:
                        return f"click[{a}]"
        return text

    def get_agents_types(self):
        return [a.type for a in self.agents]

    def prompt_layer(self, agent, available_actions):
        if agent.type == "Search_Webrun_Agent":
            return self.ask_search_agent()
        if agent.type == "Click_Webrun_Agent":
            return self.ask_click_agent()
        return 0

    def ask_search_agent(self):
        initial = self.task[self.cur_session]
        prompt = f"Instruction:\n{initial}\n"
        search_label = False
        for idx, agent_type in enumerate(self.agent_calls[self.cur_session]):
            if agent_type == "Search_Webrun_Agent":
                search_label = True
                action = self.actions[self.cur_session][idx]
                prompt += f"You have already {action}.\n"
        if search_label:
            prompt += "Can you change your search query?"
        return prompt

    def ask_click_agent(self):
        prompt = ""
        for idx in range(1, len(self.actions[self.cur_session])):
            agent_type = self.agent_calls[self.cur_session][idx]
            action = self.actions[self.cur_session][idx]
            observation = self.observations[self.cur_session][idx]
            if agent_type == "Click_Webrun_Agent":
                prompt += f"\n\nAction:{action}\nObservation:{observation}"
            else:
                prompt += f"\n\n{agent_type}:{action}\nObservation:{observation}"
        return prompt

    def call_agent(self, observation, available_actions):
        for agent in self.agents:
            if self.control_dict[agent.type] in available_actions:
                return agent

    def llm_layer(self, prompt):
        return self.llm(prompt)


    def prompt_layer(self, agent, available_actions):
        if agent.type == "Search_Webrun_Agent":
            return self.ask_search_agent()
        if agent.type == "Click_Webrun_Agent":
            return self.ask_click_agent()
        return 0

    def ask_search_agent(self):
        initial = self.task[self.cur_session]
        prompt = f"Instruction:\n{initial}\n"
        search_label = False
        for idx, agent_type in enumerate(self.agent_calls[self.cur_session]):
            if agent_type == "Search_Webrun_Agent":
                search_label = True
                action = self.actions[self.cur_session][idx]
                prompt += f"You have already {action}.\n"
        if search_label:
            prompt += "Can you change your search query?"
        return prompt

    def ask_click_agent(self):
        prompt = ""
        for idx in range(1, len(self.actions[self.cur_session])):
            agent_type = self.agent_calls[self.cur_session][idx]
            action = self.actions[self.cur_session][idx]
            observation = self.observations[self.cur_session][idx]
            if agent_type == "Click_Webrun_Agent":
                prompt += f"\n\nAction:{action}\nObservation:{observation}"
            else:
                prompt += f"\n\n{agent_type}:{action}\nObservation:{observation}"
        return prompt

    def call_agent(self, observation, available_actions):
        for agent in self.agents:
            if self.control_dict[agent.type] in available_actions:
                return agent

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, observation, available_actions=None):
        self.observations[self.cur_session].append(observation)
        agent = self.call_agent(observation, available_actions)
        prompt = self.prompt_layer(agent, available_actions)
        action = agent.forward(prompt, available_actions).lstrip()
        action = self.action_parser(action, available_actions)
        self.actions[self.cur_session].append(action)
        self.agent_calls[self.cur_session].append(agent.type)
        return action


            if agent_type == "Search_Webrun_Agent":
                search_label = True
                action = self.actions[self.cur_session][idx]
                prompt += f"You have already {action}.\n"
        if search_label:
            prompt += "Can you change your search query?"
        return prompt

    def ask_click_agent(self):
        prompt = ""
        for idx in range(1, len(self.actions[self.cur_session])):
            agent_type = self.agent_calls[self.cur_session][idx]
            action = self.actions[self.cur_session][idx]
            observation = self.observations[self.cur_session][idx]
            if agent_type == "Click_Webrun_Agent":
                prompt += f"\n\nAction:{action}\nObservation:{observation}"
            else:
                prompt += f"\n\n{agent_type}:{action}\nObservation:{observation}"
        return prompt

    def call_agent(self, observation, available_actions):
        for agent in self.agents:
            if self.control_dict[agent.type] in available_actions:
                return agent

    def llm_layer(self, prompt):
        return self.llm(prompt)

    def forward(self, observation, available_actions=None):
        self.observations[self.cur_session].append(observation)
        agent = self.call_agent(observation, available_actions)
        prompt = self.prompt_layer(agent, available_actions)
        action = agent.forward(prompt, available_actions).lstrip()
        action = self.action_parser(action, available_actions)
        self.actions[self.cur_session].append(action)
        self.agent_calls[self.cur_session].append(agent.type)
        return action


"""
 Copyright (c) 2023, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: Apache License 2.0
 For full license text, see the LICENSE file in the repo root or https://www.apache.org/licenses/LICENSE-2.0
"""

zeroshot = """Following the instruction to simulate the web shopping. Based on your observation, you may operate one of the following actions
search: using this action to search products.
click: click observed [button], the clickable buttons are in [], for example click [Back to Search] can return back to research.
Use your best knowledge and take actions. You response should follow action[detail] format, such as search[3 ounce bright citrus deodorant sensitive skin], click[B078GWRC1J], click[bright citrus].
"""

zeroshot_think = """You are simulating web shopping. Based on your observation, you may operate one of the following actions:
search: using this action to search products.
think: using this action to plan your actions and reasoning.
click: using this action to click observed [button], the clickable buttons are in [], for example click [Back to Search] can return back to research.
Use your best knowledge and take actions. You response should follow action[detail] format, such as search[3 ounce bright citrus deodorant sensitive skin], think[B078GWRC1J and B078GTKVXY are bright citrus deodorant less then 50 dollars. I can check B078GWRC1J first.], click[B078GWRC1J], click[bright citrus].
"""

oneshot = """Webshop
Instruction:
i would like a 3 ounce bottle of bright citrus deodorant for sensitive skin, and price lower than 50.00 dollars
[Search]

Action: search[3 ounce bright citrus deodorant sensitive skin]
Observation:
[Back to Search]
Page 1 (Total results: 50)
[Next >]
[B078GWRC1J]
Bright Citrus Deodorant by Earth Mama | Natural and Safe for Sensitive Skin, Pregnancy and Breastfeeding, Contains Organic Calendula 3-Ounce
$10.99
[B078GTKVXY]
Ginger Fresh Deodorant by Earth Mama | Natural and Safe for Sensitive Skin, Pregnancy and Breastfeeding, Contains Organic Calendula 3-Ounce
$10.99
[B08KBVJ4XN]
Barrel and Oak - Aluminum-Free Deodorant, Deodorant for Men, Essential Oil-Based Scent, 24-Hour Odor Protection, Cedar & Patchouli Blend, Gentle on Sensitive Skin (Mountain Sage, 2.7 oz, 2-Pack)
$15.95

Action: think[B078GWRC1J and B078GTKVXY are bright citrus deodorant less then 50 dollars. I can check B078GWRC1J first.]
Observation: OK.

Action: click[B078GWRC1J]
Observation:
[Back to Search]
[< Prev]
scent [assorted scents][bright citrus][calming lavender][ginger fresh][simply non-scents]
size [travel set (4-pack)][3 ounce (pack of 1)][3-ounce (2-pack)]
Bright Citrus Deodorant by Earth Mama | Natural and Safe for Sensitive Skin, Pregnancy and Breastfeeding, Contains Organic Calendula 3-Ounce
Price: $10.99
Rating: N.A.
[Description]
[Features]
[Reviews]
[Buy Now]

Action: think[For 3 ounce bottle of bright citrus deodorant for sensitive skin, the item has options 'bright citrus' and '3 ounce (pack of 1)' and seems good to buy.]
Observation: OK.

Action: click[bright citrus]
Observation: You have clicked bright citrus.

Action: click[3 ounce (pack of 1)]
Observation: You have clicked 3 ounce (pack of 1).

Action: click[Buy Now]
"""

plan = """Webshop
Instruction: i would like a 3 ounce bottle of bright citrus deodorant for sensitive skin, and price lower than 50.00 dollars.
Plan: First, I need to search[3 ounce bright citrus deodorant sensitive skin]. Next, I need to click the most relevant item to get more details, e.g. click[B078GWRC1J]. Next, if this item is not what I want, I will click[< Prev] to check other searching results or I can click[Back to Search] to search with different keywords. If I find one item meets the instruction requirements, I will click necessary options, such as click[bright citrus] and click[3 ounce (pack of 1)]. Finally, I will click[Buy Now] to complete the shopping.
"""

oneshot_plan = """Webshop
Instruction: i would like a 3 ounce bottle of bright citrus deodorant for sensitive skin, and price lower than 50.00 dollars.
Plan: First, I need to search[3 ounce bright citrus deodorant sensitive skin]. Next, I need to click the most relevant item to get more details, e.g. click[B078GWRC1J]. Next, if this item is not what I want, I will click[< Prev] to check other searching results or I can click[Back to Search] to search with different keywords. If I find one item meets the instruction requirements, I will click necessary options, such as click[bright citrus] and click[3 ounce (pack of 1)]. Finally, I will click[Buy Now] to complete the shopping.

Action: search[3 ounce bright citrus deodorant sensitive skin]
Observation:
[Back to Search]
Page 1 (Total results: 50)
[Next >]
[B078GWRC1J]
Bright Citrus Deodorant by Earth Mama | Natural and Safe for Sensitive Skin, Pregnancy and Breastfeeding, Contains Organic Calendula 3-Ounce
$10.99
[B078GTKVXY]
Ginger Fresh Deodorant by Earth Mama | Natural and Safe for Sensitive Skin, Pregnancy and Breastfeeding, Contains Organic Calendula 3-Ounce
$10.99
[B08KBVJ4XN]
Barrel and Oak - Aluminum-Free Deodorant, Deodorant for Men, Essential Oil-Based Scent, 24-Hour Odor Protection, Cedar & Patchouli Blend, Gentle on Sensitive Skin (Mountain Sage, 2.7 oz, 2-Pack)
$15.95

Action: click[B078GWRC1J]
Observation:
[Back to Search]
[< Prev]
scent [assorted scents][bright citrus][calming lavender][ginger fresh][simply non-scents]
size [travel set (4-pack)][3 ounce (pack of 1)][3-ounce (2-pack)]
Bright Citrus Deodorant by Earth Mama | Natural and Safe for Sensitive Skin, Pregnancy and Breastfeeding, Contains Organic Calendula 3-Ounce
Price: $10.99
Rating: N.A.
[Description]
[Features]
[Reviews]
[Buy Now]

Action: click[bright citrus]
Observation: You have clicked bright citrus.

Action: click[3 ounce (pack of 1)]
Observation: You have clicked 3 ounce (pack of 1).

Action: click[Buy Now]
"""

plan_react = """Webshop
Instruction: i would like a 3 ounce bottle of bright citrus deodorant for sensitive skin, and price lower than 50.00 dollars.
Plan: First, I need to search[3 ounce bright citrus deodorant sensitive skin]. Next, I need to think[some items are bright citrus deodorant less then 50 dollars. I can check one first.]. Next, I need to click the most relevant item to get more details, e.g. click[B078GWRC1J]. If I need think whether this item meets the requirements, e.g. think[For 3 ounce bottle of bright citrus deodorant for sensitive skin, the item has options 'bright citrus' and '3 ounce (pack of 1)' and seems good to buy.]. Then, I will click necessary options, such as click[bright citrus] and click[3 ounce (pack of 1)]. Finally, I will click[Buy Now] to complete the shopping.
"""

oneshot_plan_react = """Webshop
Instruction:
i would like a 3 ounce bottle of bright citrus deodorant for sensitive skin, and price lower than 50.00 dollars
[Search]

Action: search[3 ounce bright citrus deodorant sensitive skin]
Observation:
[Back to Search]
Page 1 (Total results: 50)
[Next >]
[B078GWRC1J]
Bright Citrus Deodorant by Earth Mama | Natural and Safe for Sensitive Skin, Pregnancy and Breastfeeding, Contains Organic Calendula 3-Ounce
$10.99
[B078GTKVXY]
Ginger Fresh Deodorant by Earth Mama | Natural and Safe for Sensitive Skin, Pregnancy and Breastfeeding, Contains Organic Calendula 3-Ounce
$10.99
[B08KBVJ4XN]
Barrel and Oak - Aluminum-Free Deodorant, Deodorant for Men, Essential Oil-Based Scent, 24-Hour Odor Protection, Cedar & Patchouli Blend, Gentle on Sensitive Skin (Mountain Sage, 2.7 oz, 2-Pack)
$15.95

Action: think[B078GWRC1J and B078GTKVXY are bright citrus deodorant less then 50 dollars. I can check B078GWRC1J first.]
Observation: OK.

Action: click[B078GWRC1J]
Observation:
[Back to Search]
[< Prev]
scent [assorted scents][bright citrus][calming lavender][ginger fresh][simply non-scents]
size [travel set (4-pack)][3 ounce (pack of 1)][3-ounce (2-pack)]
Bright Citrus Deodorant by Earth Mama | Natural and Safe for Sensitive Skin, Pregnancy and Breastfeeding, Contains Organic Calendula 3-Ounce
Price: $10.99
Rating: N.A.
[Description]
[Features]
[Reviews]
[Buy Now]

Action: think[For 3 ounce bottle of bright citrus deodorant for sensitive skin, the item has options 'bright citrus' and '3 ounce (pack of 1)' and seems good to buy.]
Observation: OK.

Action: click[bright citrus]
Observation: You have clicked bright citrus.

Action: click[3 ounce (pack of 1)]
Observation: You have clicked 3 ounce (pack of 1).

Action: click[Buy Now]
"""

search = """Webshop
Instruction:
i'm looking for a yellow hair styling product that is made from natural ingredients and easy to use, and price lower than 40.00 dollars
You have already search[yellow hair styling product made from natural ingredients easy to use price < 40.00 dollars]. Could you please change your search query?

Action: search[yellow hair styling natural ingredients]

Instruction:
i would like a 3 ounce bottle of bright citrus deodorant for sensitive skin, and price lower than 50.00 dollars

Action: search[3 ounce bright citrus deodorant sensitive skin]

Instruction:
i am looking for black pull-out organizers with nickel finish and size must be 15\" wide, and price lower than 90.00 dollars

Action: search[black pull-out organizers with nickel finish, 15\" wide]
"""

click = """Webshop
Instruction:
i would like a 3 ounce bottle of bright citrus deodorant for sensitive skin, and price lower than 50.00 dollars
[Search]

search[3 ounce bright citrus deodorant sensitive skin]
Observation:
[Back to Search]
Page 1 (Total results: 50)
[Next >]
[B078GWRC1J]
Bright Citrus Deodorant by Earth Mama | Natural and Safe for Sensitive Skin, Pregnancy and Breastfeeding, Contains Organic Calendula 3-Ounce
$10.99
[B078GTKVXY]
Ginger Fresh Deodorant by Earth Mama | Natural and Safe for Sensitive Skin, Pregnancy and Breastfeeding, Contains Organic Calendula 3-Ounce
$10.99
[B08KBVJ4XN]
Barrel and Oak - Aluminum-Free Deodorant, Deodorant for Men, Essential Oil-Based Scent, 24-Hour Odor Protection, Cedar & Patchouli Blend, Gentle on Sensitive Skin (Mountain Sage, 2.7 oz, 2-Pack)
$15.95
current available action is click:[Back to Search\tB078GWRC1J\tB078GTKVXY\tB08KBVJ4XN]

Action: click[B078GWRC1J]
Observation:
[Back to Search]
[< Prev]
scent [assorted scents][bright citrus][calming lavender][ginger fresh][simply non-scents]
size [travel set (4-pack)][3 ounce (pack of 1)][3-ounce (2-pack)]
Bright Citrus Deodorant by Earth Mama | Natural and Safe for Sensitive Skin, Pregnancy and Breastfeeding, Contains Organic Calendula 3-Ounce
Price: $10.99
Rating: N.A.
[Description]
[Features]
[Reviews]
[Buy Now]
current available action is click:[Back to Search\t< Prev\tassorted scents\tbright citrus\tcalming lavender\tginger fresh\tsimply non-scents\ttravel set (4-pack)\t3 ounce (pack of 1)\t3-ounce (2-pack)\tDescription\tFeatures\tReviews\tBuy Now]

Action: click[bright citrus]
Observation: You have clicked bright citrus.
current available action is click:[Back to Search\t< Prev\tassorted scents\tbright citrus\tcalming lavender\tginger fresh\tsimply non-scents\ttravel set (4-pack)\t3 ounce (pack of 1)\t3-ounce (2-pack)\tDescription\tFeatures\tReviews\tBuy Now]

Action: click[3 ounce (pack of 1)]
Observation: You have clicked 3 ounce (pack of 1).
current available action is click:[Back to Search\t< Prev\tassorted scents\tbright citrus\tcalming lavender\tginger fresh\tsimply non-scents\ttravel set (4-pack)\t3 ounce (pack of 1)\t3-ounce (2-pack)\tDescription\tFeatures\tReviews\tBuy Now]

Action: click[Buy Now]
"""


"""
 Copyright (c) 2023, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: Apache License 2.0
 For full license text, see the LICENSE file in the repo root or https://www.apache.org/licenses/LICENSE-2.0
"""

import json
import os
import re
from collections import defaultdict
from web_run.config import available_agent_names

def get_complexity(ins_attr_file="../webshop/data/items_human_ins_new.json"):
    with open(ins_attr_file) as f:
        human_attr = json.load(f)
    complexity_dict = {}
    for asin in human_attr:
        for ins in human_attr[asin]:
            instruction = ins["instruction"]
            attribute_complexity = len(ins["instruction_attributes"])
            option_complexity = len(ins["instruction_options"])
            key = instruction.strip('.')
            complexity_dict[key] = [attribute_complexity, option_complexity]
    return complexity_dict

def get_reward(session):
    last_session = session[-1]
    observation = last_session['observation']
    if observation.startswith('Your score'):
        tokens = observation.split(':')
        for idx, token in enumerate(tokens):
            if "Your score (min 0.0, max 1.0)" in token:
                cur_idx = idx+1
                reward = tokens[cur_idx].strip()
                return float(reward)
    return 0.0

def get_sess_ins(session):
    instruction = session[0]['instruction_text']
    return instruction

def get_sess_step(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)
    sess_idx, step_len = numbers
    return int(step_len)

def filter_results(file_name):
    sessions = get_file_sess(file_name)
    original = len(sessions)
    new_sessions = []
    with open(file_name + '.back', 'a') as f:
        for _, sess in enumerate(sessions):
            json.dump(sess, f)
            f.write('\n')
            saving_label = True
            for step in sess:
                if "action" in step: # checking whether it is a action step
                    if "No response" in step["action"]: # checking whether it is No response due to LLM
                        saving_label = False
                if "observation" in step:
                    if "Your score (min 0.0, max 1.0):" in step["observation"]:
                        saving_label = True
            if saving_label:
                new_sessions.append(sess)
    after_filtering = len(new_sessions)
    if original - after_filtering > 0:
        print(f"{original - after_filtering} session contain No response")
    with open(file_name, 'w') as f:
        for sess in new_sessions:
            json.dump(sess, f)
            f.write('\n')

def delete_repeat(file_name):
    sessions = get_file_sess(file_name)
    idx_score = {}
    idx_ref = {}
    with open(file_name + '.repeat', 'a') as f: # only saving highest score
        for sess in sessions:
            json.dump(sess, f)
            f.write('\n')
            sess_idx = get_sess_idx(sess)
            reward = get_reward(sess)
            if sess_idx in idx_score:
                if reward < idx_score[sess_idx]:
                    continue
            idx_score[sess_idx] = reward
            idx_ref[sess_idx] = sess
    with open(file_name, 'w') as f:
        sort_idx = sorted(idx_ref.keys())
        for sess_idx in sort_idx:
            sess = idx_ref[sess_idx]
            json.dump(sess, f)
            f.write('\n')

def get_file_sess(file_name):
    sessions = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            sessions.append(session)
    return sessions

def get_sess_idx(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)
    sess_idx, step_len = numbers
    return int(sess_idx)

def get_file_sess_idx(file_name):
    session_idx = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            session_idx.append(get_sess_idx(session))
    return set(session_idx)

def get_goal_item(session):
    goal = session[0]
    asin = goal['asin']
    return asin

def get_session_items(session):
    items = []
    for step in session:
        if "retrieved_items" in step:
            items += step["retrieved_items"]
    return items

def merge_sessions_by_llm(model_name, path = "./execution_data/"):
    file_list = []
    sessions = []
    for root, _, files in os.walk(path):
        for file_name in files:
            if model_name in file_name:
                file_list.append(os.path.join(root, file_name))
    for f_name in file_list:
        with open(f_name) as f:
            for line in f:
                session = json.loads(line)
                sessions.append(session)
    return sessions

def get_recall(sess):
    goal_item = get_goal_item(sess)
    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1
    else:
        return 0

def get_precision(sess):
    goal_item = get_goal_item(sess)
    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1/len(retri_items)
    else:
        return 0

def get_fail(sess):
    if "observation" in sess[-1]:
        if "Your score (min 0.0, max 1.0):" in sess[-1]["observation"]:
            return 0
    return 1

def get_halluci(sess):
    halluci_step_idx = []
    for step in sess:
        if "step_id" in step:
            step_id = step["step_id"]
            numbers = re.findall(r'\d+', step_id)
            sess_idx, step_idx = numbers
            if "observation" in step:
                if "Invalid action!" in step["observation"]:
                    halluci_step_idx.append(step_idx)
    return halluci_step_idx

def eval_complexity(sessions, complexity_dict):
    complexity = []
    for sess in sessions:
        instruction = get_sess_ins(sess)
        if instruction.find(", and price lower") > 0:
            ins_key = instruction[:instruction.find(", and price lower")]
        else:
            ins_key = instruction
        if ins_key in complexity_dict:
            attr_complexity, option_complexity = complexity_dict[ins_key]
        else:
            print("no complexity for ", ins_key)
        complexity.append(attr_complexity + option_complexity)
    return complexity

def eval_success(sessions) -> list:
    success = []
    for sess in sessions:
        r = get_reward(sess)
        if r >= 1:
            success.append(1)
        else:
            success.append(0)
    return success

def eval_initial_halluci(sessions) -> list:
    initial_halluci = []
    for sess in sessions:
        initial_halluci.append(get_halluci(sess)[0])
    return initial_halluci

def eval_avg_halluci(sessions) -> list:
    avg_halluci = []
    for sess in sessions:
        halluci = get_halluci(sess)
        avg_halluci.append(sum(halluci)/len(halluci))
    return avg_halluci

def eval_fail(sessions) -> list:
    fail_list = [get_fail(sess) for sess in sessions]
    return fail_list

def eval_reward(sessions) -> list:
    reward = []
    for sess in sessions:
        reward.append(get_reward(sess))
    return reward

def eval_recall(sessions) -> list:
    recall = [get_recall(sess) for sess in sessions]
    return recall

def eval_precision(sessions) -> list:
    prec = [get_precision(sess) for sess in sessions]
    return prec

def y_wrt_x_rel(x, y):
    x_dict = defaultdict(list)
    for idx, c in enumerate(x):
        x_dict[c].append(y[idx])
    x_dict_y = {key:round(sum(val)/len(val), 4) for key, val in x_dict.items()}
    return x_dict_y

def get_all_eval_values(llm_name):
    agent_types = available_agent_names
    for a_type in agent_types:
        get_eval_values(llm_name, a_type)

def get_eval_values(llm_name, agent_type, path="./execution_data/"):
    file_name = f"{agent_type}_{llm_name}_batch.json"
    file_name = os.path.join(path, file_name)
    delete_repeat(file_name)
    session_idx = get_file_sess_idx(file_name)
    assert len(session_idx) == 900
    sessions = get_file_sess(file_name)
    # print(len(sessions))
    recalls = eval_recall(sessions)
    rewards = eval_reward(sessions)
    precisions = eval_precision(sessions)
    success = eval_success(sessions)
    fail = eval_fail(sessions)
    avg_reward = round(sum(rewards)/len(rewards), 4)
    success_rate = round(sum(success)/len(success), 4)
    fail_rate = round(sum(fail)/len(success), 4)
    avg_recall = round(sum(recalls)/len(recalls), 4)
    avg_prec = round(sum(precisions)/len(precisions), 4)
    print(f"{avg_reward}\t{success_rate}\t{fail_rate}\t{avg_recall}\t{avg_prec}")


"""
 Copyright (c) 2023, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: Apache License 2.0
 For full license text, see the LICENSE file in the repo root or https://www.apache.org/licenses/LICENSE-2.0
"""

import json
import os
import re
from collections import defaultdict
from web_run.config import available_agent_names

def get_complexity(ins_attr_file="../webshop/data/items_human_ins_new.json"):
    with open(ins_attr_file) as f:
        human_attr = json.load(f)
    complexity_dict = {}
    for asin in human_attr:
        for ins in human_attr[asin]:
            instruction = ins["instruction"]
            attribute_complexity = len(ins["instruction_attributes"])
            option_complexity = len(ins["instruction_options"])
            key = instruction.strip('.')
            complexity_dict[key] = [attribute_complexity, option_complexity]
    return complexity_dict

def get_reward(session):
    last_session = session[-1]
    observation = last_session['observation']
    if observation.startswith('Your score'):
        tokens = observation.split(':')
        for idx, token in enumerate(tokens):
            if "Your score (min 0.0, max 1.0)" in token:
                cur_idx = idx+1
                reward = tokens[cur_idx].strip()
                return float(reward)
    return 0.0

def get_sess_ins(session):
    instruction = session[0]['instruction_text']
    return instruction

def get_sess_step(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)


import json
import os
import re
from collections import defaultdict
from web_run.config import available_agent_names

def get_complexity(ins_attr_file="../webshop/data/items_human_ins_new.json"):
    with open(ins_attr_file) as f:
        human_attr = json.load(f)
    complexity_dict = {}
    for asin in human_attr:
        for ins in human_attr[asin]:
            instruction = ins["instruction"]
            attribute_complexity = len(ins["instruction_attributes"])
            option_complexity = len(ins["instruction_options"])
            key = instruction.strip('.')
            complexity_dict[key] = [attribute_complexity, option_complexity]
    return complexity_dict

def get_reward(session):
    last_session = session[-1]
    observation = last_session['observation']
    if observation.startswith('Your score'):
        tokens = observation.split(':')
        for idx, token in enumerate(tokens):
            if "Your score (min 0.0, max 1.0)" in token:
                cur_idx = idx+1
                reward = tokens[cur_idx].strip()
                return float(reward)
    return 0.0

def get_sess_ins(session):
    instruction = session[0]['instruction_text']
    return instruction

def get_sess_step(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)
    sess_idx, step_len = numbers
    return int(step_len)

def filter_results(file_name):
    sessions = get_file_sess(file_name)
    original = len(sessions)
    new_sessions = []
    with open(file_name + '.back', 'a') as f:
        for _, sess in enumerate(sessions):
            json.dump(sess, f)
            f.write('\n')
            saving_label = True


        for ins in human_attr[asin]:
            instruction = ins["instruction"]
            attribute_complexity = len(ins["instruction_attributes"])
            option_complexity = len(ins["instruction_options"])
            key = instruction.strip('.')
            complexity_dict[key] = [attribute_complexity, option_complexity]
    return complexity_dict

def get_reward(session):
    last_session = session[-1]
    observation = last_session['observation']
    if observation.startswith('Your score'):
        tokens = observation.split(':')
        for idx, token in enumerate(tokens):
            if "Your score (min 0.0, max 1.0)" in token:
                cur_idx = idx+1
                reward = tokens[cur_idx].strip()
                return float(reward)
    return 0.0

def get_sess_ins(session):
    instruction = session[0]['instruction_text']
    return instruction

def get_sess_step(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)
    sess_idx, step_len = numbers
    return int(step_len)

def filter_results(file_name):
    sessions = get_file_sess(file_name)
    original = len(sessions)
    new_sessions = []
    with open(file_name + '.back', 'a') as f:
        for _, sess in enumerate(sessions):
            json.dump(sess, f)
            f.write('\n')
            saving_label = True
            for step in sess:
                if "action" in step: # checking whether it is a action step
                    if "No response" in step["action"]: # checking whether it is No response due to LLM
                        saving_label = False


            key = instruction.strip('.')
            complexity_dict[key] = [attribute_complexity, option_complexity]
    return complexity_dict

def get_reward(session):
    last_session = session[-1]
    observation = last_session['observation']
    if observation.startswith('Your score'):
        tokens = observation.split(':')
        for idx, token in enumerate(tokens):
            if "Your score (min 0.0, max 1.0)" in token:
                cur_idx = idx+1
                reward = tokens[cur_idx].strip()
                return float(reward)
    return 0.0

def get_sess_ins(session):
    instruction = session[0]['instruction_text']
    return instruction

def get_sess_step(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)
    sess_idx, step_len = numbers
    return int(step_len)

def filter_results(file_name):
    sessions = get_file_sess(file_name)
    original = len(sessions)
    new_sessions = []
    with open(file_name + '.back', 'a') as f:
        for _, sess in enumerate(sessions):
            json.dump(sess, f)
            f.write('\n')
            saving_label = True
            for step in sess:
                if "action" in step: # checking whether it is a action step
                    if "No response" in step["action"]: # checking whether it is No response due to LLM
                        saving_label = False
                if "observation" in step:
                    if "Your score (min 0.0, max 1.0):" in step["observation"]:
                        saving_label = True
            if saving_label:
                new_sessions.append(sess)
    after_filtering = len(new_sessions)


    observation = last_session['observation']
    if observation.startswith('Your score'):
        tokens = observation.split(':')
        for idx, token in enumerate(tokens):
            if "Your score (min 0.0, max 1.0)" in token:
                cur_idx = idx+1
                reward = tokens[cur_idx].strip()
                return float(reward)
    return 0.0

def get_sess_ins(session):
    instruction = session[0]['instruction_text']
    return instruction

def get_sess_step(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)
    sess_idx, step_len = numbers
    return int(step_len)

def filter_results(file_name):
    sessions = get_file_sess(file_name)
    original = len(sessions)
    new_sessions = []
    with open(file_name + '.back', 'a') as f:
        for _, sess in enumerate(sessions):
            json.dump(sess, f)
            f.write('\n')
            saving_label = True
            for step in sess:
                if "action" in step: # checking whether it is a action step
                    if "No response" in step["action"]: # checking whether it is No response due to LLM
                        saving_label = False
                if "observation" in step:
                    if "Your score (min 0.0, max 1.0):" in step["observation"]:
                        saving_label = True
            if saving_label:
                new_sessions.append(sess)
    after_filtering = len(new_sessions)
    if original - after_filtering > 0:
        print(f"{original - after_filtering} session contain No response")
    with open(file_name, 'w') as f:
        for sess in new_sessions:
            json.dump(sess, f)
            f.write('\n')

def delete_repeat(file_name):
    sessions = get_file_sess(file_name)
    idx_score = {}
    idx_ref = {}
    with open(file_name + '.repeat', 'a') as f: # only saving highest score
        for sess in sessions:
            json.dump(sess, f)
            f.write('\n')
            sess_idx = get_sess_idx(sess)
            reward = get_reward(sess)
            if sess_idx in idx_score:
                if reward < idx_score[sess_idx]:
                    continue
            idx_score[sess_idx] = reward
            idx_ref[sess_idx] = sess
    with open(file_name, 'w') as f:
        sort_idx = sorted(idx_ref.keys())
        for sess_idx in sort_idx:
            sess = idx_ref[sess_idx]


            json.dump(sess, f)
            f.write('\n')
            saving_label = True
            for step in sess:
                if "action" in step: # checking whether it is a action step
                    if "No response" in step["action"]: # checking whether it is No response due to LLM
                        saving_label = False
                if "observation" in step:
                    if "Your score (min 0.0, max 1.0):" in step["observation"]:
                        saving_label = True
            if saving_label:
                new_sessions.append(sess)
    after_filtering = len(new_sessions)
    if original - after_filtering > 0:
        print(f"{original - after_filtering} session contain No response")
    with open(file_name, 'w') as f:
        for sess in new_sessions:
            json.dump(sess, f)
            f.write('\n')

def delete_repeat(file_name):
    sessions = get_file_sess(file_name)
    idx_score = {}
    idx_ref = {}
    with open(file_name + '.repeat', 'a') as f: # only saving highest score
        for sess in sessions:
            json.dump(sess, f)
            f.write('\n')
            sess_idx = get_sess_idx(sess)
            reward = get_reward(sess)
            if sess_idx in idx_score:
                if reward < idx_score[sess_idx]:
                    continue
            idx_score[sess_idx] = reward
            idx_ref[sess_idx] = sess
    with open(file_name, 'w') as f:
        sort_idx = sorted(idx_ref.keys())
        for sess_idx in sort_idx:
            sess = idx_ref[sess_idx]
            json.dump(sess, f)
            f.write('\n')

def get_file_sess(file_name):
    sessions = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            sessions.append(session)
    return sessions

def get_sess_idx(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)
    sess_idx, step_len = numbers
    return int(sess_idx)

def get_file_sess_idx(file_name):
    session_idx = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)


    idx_score = {}
    idx_ref = {}
    with open(file_name + '.repeat', 'a') as f: # only saving highest score
        for sess in sessions:
            json.dump(sess, f)
            f.write('\n')
            sess_idx = get_sess_idx(sess)
            reward = get_reward(sess)
            if sess_idx in idx_score:
                if reward < idx_score[sess_idx]:
                    continue
            idx_score[sess_idx] = reward
            idx_ref[sess_idx] = sess
    with open(file_name, 'w') as f:
        sort_idx = sorted(idx_ref.keys())
        for sess_idx in sort_idx:
            sess = idx_ref[sess_idx]
            json.dump(sess, f)
            f.write('\n')

def get_file_sess(file_name):
    sessions = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            sessions.append(session)
    return sessions

def get_sess_idx(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)
    sess_idx, step_len = numbers
    return int(sess_idx)

def get_file_sess_idx(file_name):
    session_idx = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            session_idx.append(get_sess_idx(session))
    return set(session_idx)

def get_goal_item(session):
    goal = session[0]
    asin = goal['asin']
    return asin


            if sess_idx in idx_score:
                if reward < idx_score[sess_idx]:
                    continue
            idx_score[sess_idx] = reward
            idx_ref[sess_idx] = sess
    with open(file_name, 'w') as f:
        sort_idx = sorted(idx_ref.keys())
        for sess_idx in sort_idx:
            sess = idx_ref[sess_idx]
            json.dump(sess, f)
            f.write('\n')

def get_file_sess(file_name):
    sessions = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            sessions.append(session)
    return sessions

def get_sess_idx(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)
    sess_idx, step_len = numbers
    return int(sess_idx)

def get_file_sess_idx(file_name):
    session_idx = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            session_idx.append(get_sess_idx(session))
    return set(session_idx)

def get_goal_item(session):
    goal = session[0]
    asin = goal['asin']
    return asin

def get_session_items(session):
    items = []
    for step in session:
        if "retrieved_items" in step:
            items += step["retrieved_items"]
    return items


        sort_idx = sorted(idx_ref.keys())
        for sess_idx in sort_idx:
            sess = idx_ref[sess_idx]
            json.dump(sess, f)
            f.write('\n')

def get_file_sess(file_name):
    sessions = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            sessions.append(session)
    return sessions

def get_sess_idx(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)
    sess_idx, step_len = numbers
    return int(sess_idx)

def get_file_sess_idx(file_name):
    session_idx = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            session_idx.append(get_sess_idx(session))
    return set(session_idx)

def get_goal_item(session):
    goal = session[0]
    asin = goal['asin']
    return asin

def get_session_items(session):
    items = []
    for step in session:
        if "retrieved_items" in step:
            items += step["retrieved_items"]
    return items

def merge_sessions_by_llm(model_name, path = "./execution_data/"):
    file_list = []
    sessions = []
    for root, _, files in os.walk(path):
        for file_name in files:
            if model_name in file_name:
                file_list.append(os.path.join(root, file_name))


    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            sessions.append(session)
    return sessions

def get_sess_idx(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)
    sess_idx, step_len = numbers
    return int(sess_idx)

def get_file_sess_idx(file_name):
    session_idx = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            session_idx.append(get_sess_idx(session))
    return set(session_idx)

def get_goal_item(session):
    goal = session[0]
    asin = goal['asin']
    return asin

def get_session_items(session):
    items = []
    for step in session:
        if "retrieved_items" in step:
            items += step["retrieved_items"]
    return items

def merge_sessions_by_llm(model_name, path = "./execution_data/"):
    file_list = []
    sessions = []
    for root, _, files in os.walk(path):
        for file_name in files:
            if model_name in file_name:
                file_list.append(os.path.join(root, file_name))
    for f_name in file_list:
        with open(f_name) as f:
            for line in f:
                session = json.loads(line)
                sessions.append(session)


def get_sess_idx(session):
    step_id = session[-1]["step_id"]
    numbers = re.findall(r'\d+', step_id)
    sess_idx, step_len = numbers
    return int(sess_idx)

def get_file_sess_idx(file_name):
    session_idx = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            session_idx.append(get_sess_idx(session))
    return set(session_idx)

def get_goal_item(session):
    goal = session[0]
    asin = goal['asin']
    return asin

def get_session_items(session):
    items = []
    for step in session:
        if "retrieved_items" in step:
            items += step["retrieved_items"]
    return items

def merge_sessions_by_llm(model_name, path = "./execution_data/"):
    file_list = []
    sessions = []
    for root, _, files in os.walk(path):
        for file_name in files:
            if model_name in file_name:
                file_list.append(os.path.join(root, file_name))
    for f_name in file_list:
        with open(f_name) as f:
            for line in f:
                session = json.loads(line)
                sessions.append(session)
    return sessions

def get_recall(sess):
    goal_item = get_goal_item(sess)
    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1


def get_file_sess_idx(file_name):
    session_idx = []
    with open(file_name) as f:
        for line in f:
            session = json.loads(line)
            session_idx.append(get_sess_idx(session))
    return set(session_idx)

def get_goal_item(session):
    goal = session[0]
    asin = goal['asin']
    return asin

def get_session_items(session):
    items = []
    for step in session:
        if "retrieved_items" in step:
            items += step["retrieved_items"]
    return items

def merge_sessions_by_llm(model_name, path = "./execution_data/"):
    file_list = []
    sessions = []
    for root, _, files in os.walk(path):
        for file_name in files:
            if model_name in file_name:
                file_list.append(os.path.join(root, file_name))
    for f_name in file_list:
        with open(f_name) as f:
            for line in f:
                session = json.loads(line)
                sessions.append(session)
    return sessions

def get_recall(sess):
    goal_item = get_goal_item(sess)
    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1
    else:
        return 0

def get_precision(sess):
    goal_item = get_goal_item(sess)
    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1/len(retri_items)
    else:
        return 0

def get_fail(sess):
    if "observation" in sess[-1]:
        if "Your score (min 0.0, max 1.0):" in sess[-1]["observation"]:


    items = []
    for step in session:
        if "retrieved_items" in step:
            items += step["retrieved_items"]
    return items

def merge_sessions_by_llm(model_name, path = "./execution_data/"):
    file_list = []
    sessions = []
    for root, _, files in os.walk(path):
        for file_name in files:
            if model_name in file_name:
                file_list.append(os.path.join(root, file_name))
    for f_name in file_list:
        with open(f_name) as f:
            for line in f:
                session = json.loads(line)
                sessions.append(session)
    return sessions

def get_recall(sess):
    goal_item = get_goal_item(sess)
    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1
    else:
        return 0

def get_precision(sess):
    goal_item = get_goal_item(sess)
    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1/len(retri_items)
    else:
        return 0

def get_fail(sess):
    if "observation" in sess[-1]:
        if "Your score (min 0.0, max 1.0):" in sess[-1]["observation"]:
            return 0
    return 1

def get_halluci(sess):
    halluci_step_idx = []
    for step in sess:
        if "step_id" in step:
            step_id = step["step_id"]


    sessions = []
    for root, _, files in os.walk(path):
        for file_name in files:
            if model_name in file_name:
                file_list.append(os.path.join(root, file_name))
    for f_name in file_list:
        with open(f_name) as f:
            for line in f:
                session = json.loads(line)
                sessions.append(session)
    return sessions

def get_recall(sess):
    goal_item = get_goal_item(sess)
    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1
    else:
        return 0

def get_precision(sess):
    goal_item = get_goal_item(sess)
    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1/len(retri_items)
    else:
        return 0

def get_fail(sess):
    if "observation" in sess[-1]:
        if "Your score (min 0.0, max 1.0):" in sess[-1]["observation"]:
            return 0
    return 1

def get_halluci(sess):
    halluci_step_idx = []
    for step in sess:
        if "step_id" in step:
            step_id = step["step_id"]
            numbers = re.findall(r'\d+', step_id)
            sess_idx, step_idx = numbers
            if "observation" in step:
                if "Invalid action!" in step["observation"]:
                    halluci_step_idx.append(step_idx)
    return halluci_step_idx

def eval_complexity(sessions, complexity_dict):


                session = json.loads(line)
                sessions.append(session)
    return sessions

def get_recall(sess):
    goal_item = get_goal_item(sess)
    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1
    else:
        return 0

def get_precision(sess):
    goal_item = get_goal_item(sess)
    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1/len(retri_items)
    else:
        return 0

def get_fail(sess):
    if "observation" in sess[-1]:
        if "Your score (min 0.0, max 1.0):" in sess[-1]["observation"]:
            return 0
    return 1

def get_halluci(sess):
    halluci_step_idx = []
    for step in sess:
        if "step_id" in step:
            step_id = step["step_id"]
            numbers = re.findall(r'\d+', step_id)
            sess_idx, step_idx = numbers
            if "observation" in step:
                if "Invalid action!" in step["observation"]:
                    halluci_step_idx.append(step_idx)
    return halluci_step_idx

def eval_complexity(sessions, complexity_dict):
    complexity = []
    for sess in sessions:
        instruction = get_sess_ins(sess)
        if instruction.find(", and price lower") > 0:
            ins_key = instruction[:instruction.find(", and price lower")]
        else:


    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1
    else:
        return 0

def get_precision(sess):
    goal_item = get_goal_item(sess)
    retri_items = get_session_items(sess)
    if goal_item in retri_items:
        return 1/len(retri_items)
    else:
        return 0

def get_fail(sess):
    if "observation" in sess[-1]:
        if "Your score (min 0.0, max 1.0):" in sess[-1]["observation"]:
            return 0
    return 1

def get_halluci(sess):
    halluci_step_idx = []
    for step in sess:
        if "step_id" in step:
            step_id = step["step_id"]
            numbers = re.findall(r'\d+', step_id)
            sess_idx, step_idx = numbers
            if "observation" in step:
                if "Invalid action!" in step["observation"]:
                    halluci_step_idx.append(step_idx)
    return halluci_step_idx

def eval_complexity(sessions, complexity_dict):
    complexity = []
    for sess in sessions:
        instruction = get_sess_ins(sess)
        if instruction.find(", and price lower") > 0:
            ins_key = instruction[:instruction.find(", and price lower")]
        else:
            ins_key = instruction
        if ins_key in complexity_dict:
            attr_complexity, option_complexity = complexity_dict[ins_key]
        else:
            print("no complexity for ", ins_key)
        complexity.append(attr_complexity + option_complexity)
    return complexity

def eval_success(sessions) -> list:
    success = []
    for sess in sessions:
        r = get_reward(sess)


        return 0

def get_fail(sess):
    if "observation" in sess[-1]:
        if "Your score (min 0.0, max 1.0):" in sess[-1]["observation"]:
            return 0
    return 1

def get_halluci(sess):
    halluci_step_idx = []
    for step in sess:
        if "step_id" in step:
            step_id = step["step_id"]
            numbers = re.findall(r'\d+', step_id)
            sess_idx, step_idx = numbers
            if "observation" in step:
                if "Invalid action!" in step["observation"]:
                    halluci_step_idx.append(step_idx)
    return halluci_step_idx

def eval_complexity(sessions, complexity_dict):
    complexity = []
    for sess in sessions:
        instruction = get_sess_ins(sess)
        if instruction.find(", and price lower") > 0:
            ins_key = instruction[:instruction.find(", and price lower")]
        else:
            ins_key = instruction
        if ins_key in complexity_dict:
            attr_complexity, option_complexity = complexity_dict[ins_key]
        else:
            print("no complexity for ", ins_key)
        complexity.append(attr_complexity + option_complexity)
    return complexity

def eval_success(sessions) -> list:
    success = []
    for sess in sessions:
        r = get_reward(sess)
        if r >= 1:
            success.append(1)
        else:
            success.append(0)
    return success

def eval_initial_halluci(sessions) -> list:
    initial_halluci = []
    for sess in sessions:
        initial_halluci.append(get_halluci(sess)[0])
    return initial_halluci

def eval_avg_halluci(sessions) -> list:
    avg_halluci = []
    for sess in sessions:


            if "observation" in step:
                if "Invalid action!" in step["observation"]:
                    halluci_step_idx.append(step_idx)
    return halluci_step_idx

def eval_complexity(sessions, complexity_dict):
    complexity = []
    for sess in sessions:
        instruction = get_sess_ins(sess)
        if instruction.find(", and price lower") > 0:
            ins_key = instruction[:instruction.find(", and price lower")]
        else:
            ins_key = instruction
        if ins_key in complexity_dict:
            attr_complexity, option_complexity = complexity_dict[ins_key]
        else:
            print("no complexity for ", ins_key)
        complexity.append(attr_complexity + option_complexity)
    return complexity

def eval_success(sessions) -> list:
    success = []
    for sess in sessions:
        r = get_reward(sess)
        if r >= 1:
            success.append(1)
        else:
            success.append(0)
    return success

def eval_initial_halluci(sessions) -> list:
    initial_halluci = []
    for sess in sessions:
        initial_halluci.append(get_halluci(sess)[0])
    return initial_halluci

def eval_avg_halluci(sessions) -> list:
    avg_halluci = []
    for sess in sessions:
        halluci = get_halluci(sess)
        avg_halluci.append(sum(halluci)/len(halluci))
    return avg_halluci

def eval_fail(sessions) -> list:
    fail_list = [get_fail(sess) for sess in sessions]
    return fail_list

def eval_reward(sessions) -> list:
    reward = []


            ins_key = instruction[:instruction.find(", and price lower")]
        else:
            ins_key = instruction
        if ins_key in complexity_dict:
            attr_complexity, option_complexity = complexity_dict[ins_key]
        else:
            print("no complexity for ", ins_key)
        complexity.append(attr_complexity + option_complexity)
    return complexity

def eval_success(sessions) -> list:
    success = []
    for sess in sessions:
        r = get_reward(sess)
        if r >= 1:
            success.append(1)
        else:
            success.append(0)
    return success

def eval_initial_halluci(sessions) -> list:
    initial_halluci = []
    for sess in sessions:
        initial_halluci.append(get_halluci(sess)[0])
    return initial_halluci

def eval_avg_halluci(sessions) -> list:
    avg_halluci = []
    for sess in sessions:
        halluci = get_halluci(sess)
        avg_halluci.append(sum(halluci)/len(halluci))
    return avg_halluci

def eval_fail(sessions) -> list:
    fail_list = [get_fail(sess) for sess in sessions]
    return fail_list

def eval_reward(sessions) -> list:
    reward = []
    for sess in sessions:
        reward.append(get_reward(sess))
    return reward

def eval_recall(sessions) -> list:
    recall = [get_recall(sess) for sess in sessions]


            print("no complexity for ", ins_key)
        complexity.append(attr_complexity + option_complexity)
    return complexity

def eval_success(sessions) -> list:
    success = []
    for sess in sessions:
        r = get_reward(sess)
        if r >= 1:
            success.append(1)
        else:
            success.append(0)
    return success

def eval_initial_halluci(sessions) -> list:
    initial_halluci = []
    for sess in sessions:
        initial_halluci.append(get_halluci(sess)[0])
    return initial_halluci

def eval_avg_halluci(sessions) -> list:
    avg_halluci = []
    for sess in sessions:
        halluci = get_halluci(sess)
        avg_halluci.append(sum(halluci)/len(halluci))
    return avg_halluci

def eval_fail(sessions) -> list:
    fail_list = [get_fail(sess) for sess in sessions]
    return fail_list

def eval_reward(sessions) -> list:
    reward = []
    for sess in sessions:
        reward.append(get_reward(sess))
    return reward

def eval_recall(sessions) -> list:
    recall = [get_recall(sess) for sess in sessions]
    return recall

def eval_precision(sessions) -> list:
    prec = [get_precision(sess) for sess in sessions]
    return prec

def y_wrt_x_rel(x, y):


        r = get_reward(sess)
        if r >= 1:
            success.append(1)
        else:
            success.append(0)
    return success

def eval_initial_halluci(sessions) -> list:
    initial_halluci = []
    for sess in sessions:
        initial_halluci.append(get_halluci(sess)[0])
    return initial_halluci

def eval_avg_halluci(sessions) -> list:
    avg_halluci = []
    for sess in sessions:
        halluci = get_halluci(sess)
        avg_halluci.append(sum(halluci)/len(halluci))
    return avg_halluci

def eval_fail(sessions) -> list:
    fail_list = [get_fail(sess) for sess in sessions]
    return fail_list

def eval_reward(sessions) -> list:
    reward = []
    for sess in sessions:
        reward.append(get_reward(sess))
    return reward

def eval_recall(sessions) -> list:
    recall = [get_recall(sess) for sess in sessions]
    return recall

def eval_precision(sessions) -> list:
    prec = [get_precision(sess) for sess in sessions]
    return prec

def y_wrt_x_rel(x, y):
    x_dict = defaultdict(list)
    for idx, c in enumerate(x):
        x_dict[c].append(y[idx])
    x_dict_y = {key:round(sum(val)/len(val), 4) for key, val in x_dict.items()}