"""
 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
"""

from langchain.prompts import PromptTemplate

ZEROSHOT_INSTRUCTION = """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types:
(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search. For example, Search[Milhouse]
(2) Lookup[keyword], which returns the next sentence containing keyword in the last passage successfully found by Search. For example, Lookup[named after]
(3) Finish[answer], which returns the answer and finishes the task. For example, Finish[Richard Nixon]
You may take as many steps as necessary.
Question: {question}{scratchpad}"""

zeroshot_agent_prompt = PromptTemplate(
                        input_variables=["question", "scratchpad"],
                        template = ZEROSHOT_INSTRUCTION,
                        )

REACT_INSTRUCTION = """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types:
(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search.
(2) Lookup[keyword], which returns the next sentence containing keyword in the last passage successfully found by Search.
(3) Finish[answer], which returns the answer and finishes the task.
You may take as many steps as necessary.
Here are some examples:
{examples}
(END OF EXAMPLES)
Question: {question}{scratchpad}"""

react_agent_prompt = PromptTemplate(
                        input_variables=["examples", "question", "scratchpad"],
                        template = REACT_INSTRUCTION,
                        )

PLAN_INSTRUCTION = """Setup a plan for answering question with Actions. Action can be three types:
(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search.
(2) Lookup[keyword], which returns the next sentence containing keyword in the last passage successfully found by Search.
(3) Finish[answer], which returns the answer and finishes the task.
{examples}
(END OF EXAMPLES)
Question: {question}
Plan:"""

plan_prompt = PromptTemplate(
                input_variables=["examples", "question"],
                template = PLAN_INSTRUCTION,
            )

PLANNER_INSTRUCTION = """Solve a question answering task with Plan, interleaving Action, Observation steps. Plan is decided ahead of Actions. Action can be three types:
(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search.
(2) Lookup[keyword], which returns the next sentence containing keyword in the last passage successfully found by Search.
(3) Finish[answer], which returns the answer and finishes the task.
You may take as many steps as necessary.
Here are some examples:
{examples}
(END OF EXAMPLES)
Question: {question}
Plan: {plan}{scratchpad}"""

planner_agent_prompt = PromptTemplate(
                        input_variables=["examples", "question", "plan", "scratchpad"],
                        template = PLANNER_INSTRUCTION,
                        )

PLANNERREACT_INSTRUCTION = """Solve a question answering task with Plan, interleaving Thought, Action, Observation steps. Plan is decided ahead of Actions. Thought can reason about the current situation. Action can be three types:
(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search.
(2) Lookup[keyword], which returns the next sentence containing keyword in the last passage successfully found by Search.
(3) Finish[answer], which returns the answer and finishes the task.
You may take as many steps as necessary.
Here are some examples:
{examples}
(END OF EXAMPLES)
Question: {question}
Plan: {plan}{scratchpad}"""

plannerreact_agent_prompt = PromptTemplate(
                        input_variables=["examples", "question", "plan", "scratchpad"],
                        template = PLANNERREACT_INSTRUCTION,
                        )


"""
 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
"""

REACT_EXAMPLE = """Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who?
Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after.
Action 1: Search[Milhouse]
Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening.
Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after".
Action 2: Lookup[named after]
Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous.
Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon.
Action 3: Finish[Richard Nixon]

Question: Were Pavel Urysohn and Leonid Levin known for the same type of work?
Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same.
Action 1: Search[Pavel Urysohn]
Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory.
Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work.
Action 2: Search[Leonid Levin]
Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist.
Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work.
Action 3: Finish[yes]
"""

PLAN_EXAMPLE = """Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who?
Plan: I need to first need to Search[Milhouse] to find who it is named after with Lookup[named after]. Then I will collect the information to Finish[answer].

Question: Were Pavel Urysohn and Leonid Levin known for the same type of work?
Plan: I need to first Search[Pavel Urysohn] and then Search[Leonid Levin] to get more information. Then I will check their work type. If they are known for the same type of work, I will Finish[yes]. Otherwise I will Finish[no].

Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture?
Plan: I need to Search[Adam Clayton Powell] and Search[The Saimaa Gesture] to collect their information. If I find the answer, I will Finish[Adam Clayton Powell] or Finish[The Saimaa Gesture]. If I cannot find the answer, I will Lookup[rock groups] in both results.
"""

PLANNER_EXAMPLE = """Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who?
Plan: I need to first need to Search[Milhouse] to find who it is named after with Lookup[named after]. Then I will collect the information to Finish[answer].
Action 1: Search[Milhouse]
Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening.
Action 2: Lookup[named after]
Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous.
Action 3: Finish[Richard Nixon]

Question: Were Pavel Urysohn and Leonid Levin known for the same type of work?
Plan: I need to first Search[Pavel Urysohn] and then Search[Leonid Levin] to get more information. Then I will check their work type. If they are known for the same type of work, I will Finish[yes]. Otherwise I will Finish[no].
Action 1: Search[Pavel Urysohn]
Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory.
Action 2: Search[Leonid Levin]
Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist.
Action 3: Finish[yes]
"""

PLANNERREACT_EXAMPLE = """Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who?
Plan: I need to first need to Search[Milhouse] to find who it is named after with Lookup[named after]. Then I will collect the information to Finish[answer].
Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after.
Action 1: Search[Milhouse]
Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening.
Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after".
Action 2: Lookup[named after]
Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous.
Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon.
Action 3: Finish[Richard Nixon]

Question: Were Pavel Urysohn and Leonid Levin known for the same type of work?
Plan: I need to first Search[Pavel Urysohn] and then Search[Leonid Levin] to get more information. Then I will check their work type. If they are known for the same type of work, I will Finish[yes]. Otherwise I will Finish[no].
Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same.
Action 1: Search[Pavel Urysohn]
Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory.
Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work.
Action 2: Search[Leonid Levin]
Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist.
Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work.
Action 3: Finish[yes]
"""

WEBTHINK_SIMPLE6 = """Question: What is the elevation range for the area that the eastern sector of the Colorado orogeny extends into?
Thought 1: I need to search Colorado orogeny, find the area that the eastern sector of the Colorado orogeny extends into, then find the elevation range of the area.
Action 1: Search[Colorado orogeny]
Observation 1: The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas.
Thought 2: It does not mention the eastern sector. So I need to look up eastern sector.
Action 2: Lookup[eastern sector]
Observation 2: (Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny.
Thought 3: The eastern sector of Colorado orogeny extends into the High Plains. So I need to search High Plains and find its elevation range.
Action 3: Search[High Plains]
Observation 3: High Plains refers to one of two distinct land regions:
Thought 4: I need to instead search High Plains (United States).
Action 4: Search[High Plains (United States)]
Observation 4: The High Plains are a subregion of the Great Plains. From east to west, the High Plains rise in elevation from around 1,800 to 7,000 ft (550 to 2,130 m).[3]
Thought 5: High Plains rise in elevation from around 1,800 to 7,000 ft, so the answer is 1,800 to 7,000 ft.
Action 5: Finish[1,800 to 7,000 ft]

Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who?
Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after.
Action 1: Search[Milhouse]
Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening.
Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after".
Action 2: Lookup[named after]
Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous.
Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon.
Action 3: Finish[Richard Nixon]

Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture?
Thought 1: I need to search Adam Clayton Powell and The Saimaa Gesture, and find which documentary is about Finnish rock groups.
Action 1: Search[Adam Clayton Powell]
Observation 1: Could not find [Adam Clayton Powell]. Similar: ['Adam Clayton Powell III', 'Seventh Avenue (Manhattan)', 'Adam Clayton Powell Jr. State Office Building', 'Isabel Washington Powell', 'Adam Powell', 'Adam Clayton Powell (film)', 'Giancarlo Esposito'].
Thought 2: To find the documentary, I can search Adam Clayton Powell (film).
Action 2: Search[Adam Clayton Powell (film)]
Observation 2: Adam Clayton Powell is a 1989 American documentary film directed by Richard Kilberg.
The film is about the rise and fall of influential African-American politician Adam Clayton Powell Jr.[3][4] It was later aired as part of the PBS series The American Experience.
Thought 3: Adam Clayton Powell (film) is a documentary about an African-American politician, not Finnish rock groups. So the documentary about Finnish rock groups must instead be The Saimaa Gesture.
Action 3: Finish[The Saimaa Gesture]

Question: What profession does Nicholas Ray and Elia Kazan have in common?
Thought 1: I need to search Nicholas Ray and Elia Kazan, find their professions, then find the profession they have in common.
Action 1: Search[Nicholas Ray]
Observation 1: Nicholas Ray (born Raymond Nicholas Kienzle Jr., August 7, 1911 – June 16, 1979) was an American film director, screenwriter, and actor best known for the 1955 film Rebel Without a Cause.
Thought 2: Professions of Nicholas Ray are director, screenwriter, and actor. I need to search Elia Kazan next and find his professions.
Action 2: Search[Elia Kazan]
Observation 2: Elia Kazan was an American film and theatre director, producer, screenwriter and actor.
Thought 3: Professions of Elia Kazan are director, producer, screenwriter, and actor. So profession Nicholas Ray and Elia Kazan have in common is director, screenwriter, and actor.
Action 3: Finish[director, screenwriter, actor]

Question: Which magazine was started first Arthur's Magazine or First for Women?
Thought 1: I need to search Arthur's Magazine and First for Women, and find which was started first.
Action 1: Search[Arthur's Magazine]
Observation 1: Arthur's Magazine (1844-1846) was an American literary periodical published in Philadelphia in the 19th century.
Thought 2: Arthur's Magazine was started in 1844. I need to search First for Women next.
Action 2: Search[First for Women]
Observation 2: First for Women is a woman's magazine published by Bauer Media Group in the USA.[1] The magazine was started in 1989.
Thought 3: First for Women was started in 1989. 1844 (Arthur's Magazine) < 1989 (First for Women), so Arthur's Magazine was started first.
Action 3: Finish[Arthur's Magazine]

Question: Were Pavel Urysohn and Leonid Levin known for the same type of work?
Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same.
Action 1: Search[Pavel Urysohn]
Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory.
Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work.
Action 2: Search[Leonid Levin]
Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist.
Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work.
Action 3: Finish[yes]"""

REFLECTIONS = """
Previous Trial:
Question: The Rome Protocols were signed by three Prime Ministers one of which was assassinated as part of what?
Thought 1: I need to search Rome Protocols, find the three Prime Ministers, then find what they were assassinated as part of.
Action 1: Search[Rome Protocols]
Observation 1: The Rome Protocols were a series of three international agreements concluded in Rome on 17 March 1934 between the governments of Austria, Hungary and Italy. They were signed by Italian Prime Minister Benito Mussolini, Austrian Chancellor Engelbert Dollfuss and Hungarian Prime Minister Gyula Gömbös. All of the protocols went into effect on 12 July 1934 and were registered in League of Nations Treaty Series on 12 December 1934.
Thought 2: The three Prime Ministers were Benito Mussolini, Engelbert Dollfuss, and Gyula Gömbös. I need to search Benito Mussolini and find what he was assassinated as part of.
Action 2: Search[Benito Mussolini]
Observation 2: Benito Amilcare Andrea Mussolini (UK:  MU(U)SS-ə-LEE-nee, US:  MOOSS-, Italian: [beˈniːto aˈmilkare anˈdrɛːa mussoˈliːni]; 29 July 1883 – 28 April 1945) was an Italian politician and journalist who founded and led the National Fascist Party (PNF). He was Prime Minister of Italy from the March on Rome in 1922 until his deposition in 1943, as well as "Duce" of Italian fascism from the establishment of the Italian Fasces of Combat in 1919 until his summary execution in 1945 by Italian partisans. As dictator of Italy and principal founder of fascism, Mussolini inspired and supported the international spread of fascist movements during the inter-war period.Mussolini was originally a socialist politician and a journalist at the Avanti! newspaper. In 1912, he became a member of the National Directorate of the Italian Socialist Party (PSI), but he was expelled from the PSI for advocating military intervention in World War I, in opposition to the party's stance on neutrality. In 1914, Mussolini founded a new journal, Il Popolo d'Italia, and served in the Royal Italian Army during the war until he was wounded and discharged in 1917. Mussolini denounced the PSI, his views now centering on Italian nationalism instead of socialism, and later founded the fascist movement which came to oppose egalitarianism and class conflict, instead advocating "revolutionary nationalism" transcending class lines. On 31 October 1922, following the March on Rome (28–30 October), Mussolini was appointed prime minister by King Victor Emmanuel III, becoming the youngest individual to hold the office up to that time. After removing all political opposition through his secret police and outlawing labor strikes, Mussolini and his followers consolidated power through a series of laws that transformed the nation into a one-party dictatorship. Within five years, Mussolini had established dictatorial authority by both legal and illegal means and aspired to create a totalitarian state. In 1929, Mussolini signed the Lateran Treaty with the Holy See to establish Vatican City.
Mussolini's foreign policy aimed to restore the ancient grandeur of the Roman Empire by expanding Italian colonial possessions and the fascist sphere of influence. In the 1920s, he ordered the Pacification of Libya, instructed the bombing of Corfu over an incident with Greece, established a protectorate over Albania, and incorporated the city of Fiume into the Italian state via agreements with Yugoslavia. In 1936, Ethiopia was conquered following the Second Italo-Ethiopian War and merged into Italian East Africa (AOI) with Eritrea and Somalia. In 1939, Italian forces annexed Albania. Between 1936 and 1939, Mussolini ordered the successful Italian military intervention in Spain in favor of Francisco Franco during the Spanish Civil War. Mussolini's Italy initially tried to avoid the outbreak of a second global war, sending troops at the Brenner Pass to delay Anschluss and taking part in the Stresa Front, the Lytton Report, the Treaty of Lausanne, the Four-Power Pact and the Munich Agreement. However, Italy then alienated itself from Britain and France by aligning with Germany and Japan. Germany invaded Poland on 1 September 1939, resulting in declarations of war by France and the UK and the start of World War II.
On 10 June 1940, Mussolini decided to enter the war on the Axis side. Despite initial success, the subsequent Axis collapse on multiple fronts and eventual Allied invasion of Sicily made Mussolini lose the support of the population and members of the Fascist Party. As a consequence, early on 25 July 1943, the Grand Council of Fascism passed a motion of no confidence in Mussolini; later that day King Victor Emmanuel III dismissed him as head of government and had him placed in custody, appointing Pietro Badoglio to succeed him as Prime Minister. After the king agreed to an armistice with the Allies, on 12 September 1943 Mussolini was rescued from captivity in the Gran Sasso raid by German paratroopers and Waffen-SS commandos led by Major Otto-Harald Mors. Adolf Hitler, after meeting with the rescued former dictator, then put Mussolini in charge of a puppet regime in northern Italy, the Italian Social Republic (Italian: Repubblica Sociale Italiana, RSI), informally known as the Salò Republic, causing a civil war. In late April 1945, in the wake of near total defeat, Mussolini and his mistress Clara Petacci attempted to flee to Switzerland, but both were captured by Italian communist partisans and summarily executed by firing squad on 28 April 1945 near Lake Como. The bodies of Mussolini and his mistress were then taken to Milan, where they were hung upside down at a service station to publicly confirm their demise.
Thought 3: Benito Mussolini was assassinated as part of the collapse of the Axis powers during World War II. So the Rome Protocols were signed by three Prime Ministers one of which was assassinated as part of World War II.
Action 3: Finish[World War II]

Reflection: I searched one of the prime ministers involved in the signing, then attemted to answer right away. I should have searched each of the prime ministers, then looked up 'death' on each of their pages in order to get more information before answering.

Previous Trial:
Question: Kam Heskin plays Paige Morgan in a 2004 film directed by who?
Thought 1: I need to search Kam Heskin and Paige Morgan, find the 2004 film they are in, then find the director of the film.
Action 1: Search[Kam Heskin]
Observation 1: Kam Erika Heskin (born May 8, 1973) is an American actress. She began her career playing Caitlin Deschanel on the NBC daytime soap opera Sunset Beach (1998–1999), before appearing in films Planet of the Apes (2001) and Catch Me If You Can (2002). Heskin went on in 2003 to play Elizabeth Bennet in an independent film Pride & Prejudice: A Latter-Day Comedy, Robin in Charmed (2004),  and Paige Morgan in The Prince and Me film franchise (2006–2010).
Thought 2: Kam Heskin plays Paige Morgan in The Prince and Me film franchise. I need to search The Prince and Me film franchise and find the director of the 2004 film.
Action 2: Search[The Prince and Me film franchise]
Observation 2: Could not find [The Prince and Me film franchise]. Similar: ['Kam Heskin', 'Prince of Persia', 'Harry Potter and the Half-Blood Prince (film)', 'Prince of Persia: The Sands of Time (film)', 'The Chronicles of Narnia (film series)', 'Majestic Prince (franchise)', 'Kingsman (franchise)', 'The Little Prince (2015 film)', 'Shrek (franchise)', 'Harry Potter (film series)']
Thought 3: I need to search The Prince and Me (2004 film) instead.
Action 3: Search[The Prince and Me (2004 film)]
Observation 3: Could not find [The Prince and Me (2004 film)].
# ... truncated ...


"""
 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
import pandas as pd
import hotpotqa_run.utils as utils

def eval_success(result_file) -> list:
    df = pd.read_csv(result_file)
    return df['success'].tolist()

def eval_reward(result_file) -> list:
    df = pd.read_csv(result_file)
    return df['reward'].tolist()

def eval_llm_agent(llm_name, agent_name):
    levels = ['easy','medium','hard']
    all_reward = []
    all_success = []
    for l in levels:
        file_name = f"execution_data/hotpotqa/{l}_{agent_name}_{llm_name}.csv"
        all_reward += eval_reward(file_name)
        all_success += eval_success(file_name)
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_llm_agent_level(llm_name, agent_name, level):
    file_name = f"execution_data/hotpotqa/{level}_{agent_name}_{llm_name}.csv"
    all_reward = eval_reward(file_name)
    all_success = eval_success(file_name)
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_sessions(llm_name, agent_name):
    levels = ['easy','medium','hard']
    all_reward = []
    all_success = []
    for l in levels:
        reward, success = eval_sessions_level((llm_name, agent_name,l))
        all_reward += reward
        all_success += success
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_sessions_level(llm_name, agent_name,level):
    file_name = f"execution_data/hotpotqa/{level}_{agent_name}_{llm_name}.jsonl"
    sessions = utils.get_all_agent_sessions(file_name)
    all_reward = [sess["reward"] for sess in sessions]
    all_success = [sess["correct"] for sess in sessions]
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def get_reward_w_level(llm_name, agent_name):
    levels = ['easy','medium','hard']
    ret = []
    for l in levels:
        reward, _ = eval_sessions_level(llm_name, agent_name, l)
        ret.append(reward)
    return ret


"""
 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
import pandas as pd
import hotpotqa_run.utils as utils

def eval_success(result_file) -> list:
    df = pd.read_csv(result_file)
    return df['success'].tolist()

def eval_reward(result_file) -> list:
    df = pd.read_csv(result_file)
    return df['reward'].tolist()

def eval_llm_agent(llm_name, agent_name):
    levels = ['easy','medium','hard']
    all_reward = []
    all_success = []
    for l in levels:
        file_name = f"execution_data/hotpotqa/{l}_{agent_name}_{llm_name}.csv"
        all_reward += eval_reward(file_name)
        all_success += eval_success(file_name)
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_llm_agent_level(llm_name, agent_name, level):
    file_name = f"execution_data/hotpotqa/{level}_{agent_name}_{llm_name}.csv"
    all_reward = eval_reward(file_name)


"""
 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
import pandas as pd
import hotpotqa_run.utils as utils

def eval_success(result_file) -> list:
    df = pd.read_csv(result_file)
    return df['success'].tolist()

def eval_reward(result_file) -> list:
    df = pd.read_csv(result_file)
    return df['reward'].tolist()

def eval_llm_agent(llm_name, agent_name):
    levels = ['easy','medium','hard']
    all_reward = []
    all_success = []
    for l in levels:
        file_name = f"execution_data/hotpotqa/{l}_{agent_name}_{llm_name}.csv"
        all_reward += eval_reward(file_name)
        all_success += eval_success(file_name)
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_llm_agent_level(llm_name, agent_name, level):
    file_name = f"execution_data/hotpotqa/{level}_{agent_name}_{llm_name}.csv"
    all_reward = eval_reward(file_name)
    all_success = eval_success(file_name)
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success


 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
import pandas as pd
import hotpotqa_run.utils as utils

def eval_success(result_file) -> list:
    df = pd.read_csv(result_file)
    return df['success'].tolist()

def eval_reward(result_file) -> list:
    df = pd.read_csv(result_file)
    return df['reward'].tolist()

def eval_llm_agent(llm_name, agent_name):
    levels = ['easy','medium','hard']
    all_reward = []
    all_success = []
    for l in levels:
        file_name = f"execution_data/hotpotqa/{l}_{agent_name}_{llm_name}.csv"
        all_reward += eval_reward(file_name)
        all_success += eval_success(file_name)
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_llm_agent_level(llm_name, agent_name, level):
    file_name = f"execution_data/hotpotqa/{level}_{agent_name}_{llm_name}.csv"
    all_reward = eval_reward(file_name)
    all_success = eval_success(file_name)
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_sessions(llm_name, agent_name):
    levels = ['easy','medium','hard']
    all_reward = []
    all_success = []
    for l in levels:
        reward, success = eval_sessions_level((llm_name, agent_name,l))
        all_reward += reward
        all_success += success
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success


def eval_success(result_file) -> list:
    df = pd.read_csv(result_file)
    return df['success'].tolist()

def eval_reward(result_file) -> list:
    df = pd.read_csv(result_file)
    return df['reward'].tolist()

def eval_llm_agent(llm_name, agent_name):
    levels = ['easy','medium','hard']
    all_reward = []
    all_success = []
    for l in levels:
        file_name = f"execution_data/hotpotqa/{l}_{agent_name}_{llm_name}.csv"
        all_reward += eval_reward(file_name)
        all_success += eval_success(file_name)
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_llm_agent_level(llm_name, agent_name, level):
    file_name = f"execution_data/hotpotqa/{level}_{agent_name}_{llm_name}.csv"
    all_reward = eval_reward(file_name)
    all_success = eval_success(file_name)
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_sessions(llm_name, agent_name):
    levels = ['easy','medium','hard']
    all_reward = []
    all_success = []
    for l in levels:
        reward, success = eval_sessions_level((llm_name, agent_name,l))
        all_reward += reward
        all_success += success
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_sessions_level(llm_name, agent_name,level):
    file_name = f"execution_data/hotpotqa/{level}_{agent_name}_{llm_name}.jsonl"
    sessions = utils.get_all_agent_sessions(file_name)
    all_reward = [sess["reward"] for sess in sessions]
    all_success = [sess["correct"] for sess in sessions]
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)


def eval_llm_agent(llm_name, agent_name):
    levels = ['easy','medium','hard']
    all_reward = []
    all_success = []
    for l in levels:
        file_name = f"execution_data/hotpotqa/{l}_{agent_name}_{llm_name}.csv"
        all_reward += eval_reward(file_name)
        all_success += eval_success(file_name)
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_llm_agent_level(llm_name, agent_name, level):
    file_name = f"execution_data/hotpotqa/{level}_{agent_name}_{llm_name}.csv"
    all_reward = eval_reward(file_name)
    all_success = eval_success(file_name)
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_sessions(llm_name, agent_name):
    levels = ['easy','medium','hard']
    all_reward = []
    all_success = []
    for l in levels:
        reward, success = eval_sessions_level((llm_name, agent_name,l))
        all_reward += reward
        all_success += success
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_sessions_level(llm_name, agent_name,level):
    file_name = f"execution_data/hotpotqa/{level}_{agent_name}_{llm_name}.jsonl"
    sessions = utils.get_all_agent_sessions(file_name)
    all_reward = [sess["reward"] for sess in sessions]
    all_success = [sess["correct"] for sess in sessions]
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def get_reward_w_level(llm_name, agent_name):
    levels = ['easy','medium','hard']
    ret = []
    for l in levels:
        reward, _ = eval_sessions_level(llm_name, agent_name, l)
        ret.append(reward)
    return ret


def eval_llm_agent_level(llm_name, agent_name, level):
    file_name = f"execution_data/hotpotqa/{level}_{agent_name}_{llm_name}.csv"
    all_reward = eval_reward(file_name)
    all_success = eval_success(file_name)
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_sessions(llm_name, agent_name):
    levels = ['easy','medium','hard']
    all_reward = []
    all_success = []
    for l in levels:
        reward, success = eval_sessions_level((llm_name, agent_name,l))
        all_reward += reward
        all_success += success
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_sessions_level(llm_name, agent_name,level):
    file_name = f"execution_data/hotpotqa/{level}_{agent_name}_{llm_name}.jsonl"
    sessions = utils.get_all_agent_sessions(file_name)
    all_reward = [sess["reward"] for sess in sessions]
    all_success = [sess["correct"] for sess in sessions]
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def get_reward_w_level(llm_name, agent_name):
    levels = ['easy','medium','hard']
    ret = []
    for l in levels:
        reward, _ = eval_sessions_level(llm_name, agent_name, l)
        ret.append(reward)
    return ret


    levels = ['easy','medium','hard']
    all_reward = []
    all_success = []
    for l in levels:
        reward, success = eval_sessions_level((llm_name, agent_name,l))
        all_reward += reward
        all_success += success
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def eval_sessions_level(llm_name, agent_name,level):
    file_name = f"execution_data/hotpotqa/{level}_{agent_name}_{llm_name}.jsonl"
    sessions = utils.get_all_agent_sessions(file_name)
    all_reward = [sess["reward"] for sess in sessions]
    all_success = [sess["correct"] for sess in sessions]
    avg_reward = sum(all_reward)/len(all_reward)
    avg_success = sum(all_success)/len(all_success)
    return avg_reward, avg_success

def get_reward_w_level(llm_name, agent_name):
    levels = ['easy','medium','hard']
    ret = []
    for l in levels:
        reward, _ = eval_sessions_level(llm_name, agent_name, l)
        ret.append(reward)
    return ret


"""
 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 os
import sys
import json
import random
import tiktoken
token_enc = tiktoken.get_encoding("cl100k_base")
import openai
from langchain import PromptTemplate, OpenAI, LLMChain
from transformers import AutoTokenizer, AutoModelForCausalLM
from langchain.chat_models import ChatOpenAI
from langchain.chains import LLMChain
from langchain import PromptTemplate
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from web_run.config import OPENAI_API_KEY

OPENAI_CHAT_MODELS = ["gpt-3.5-turbo","gpt-3.5-turbo-16k-0613","gpt-3.5-turbo-16k","gpt-4-0613","gpt-4-32k-0613"]
OPENAI_LLM_MODELS = ["text-davinci-003","text-ada-001"]
FASTCHAT_MODELS = ["vicuna-7b"]

class langchain_openai_chatllm:
    def __init__(self, llm_name):
        openai.api_key = OPENAI_API_KEY
        self.llm_name = llm_name
        human_template="{prompt}"
        human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
        self.chat_prompt = ChatPromptTemplate.from_messages([human_message_prompt])

    def run(self, prompt, temperature=1, stop=['\n'], max_tokens=128):
        chat = ChatOpenAI(model=self.llm_name, temperature=temperature, stop=stop, max_tokens=max_tokens)
        self.chain = LLMChain(llm=chat, prompt=self.chat_prompt)
        return self.chain.run(prompt)

class langchain_openai_llm:
    def __init__(self, llm_name):
        openai.api_key = OPENAI_API_KEY
        self.prompt_temp = PromptTemplate(
            input_variables=["prompt"], template="{prompt}"
        )
        self.llm_name = llm_name

    def run(self, prompt, temperature=0.9, stop=['\n'], max_tokens=128):
        llm = OpenAI(model=self.llm_name, temperature=temperature, stop=['\n'], max_tokens=max_tokens)
        chain = LLMChain(llm=llm, prompt=self.prompt_temp)
        return chain.run(prompt)

class langchain_fastchat_llm:
    def __init__(self, llm_name):
        openai.api_key = "EMPTY" # Not support yet
        openai.api_base = "http://localhost:8000/v1"
        self.prompt_temp = PromptTemplate(
            input_variables=["prompt"], template="{prompt}"
        )
        self.llm_name = llm_name

    def run(self, prompt, temperature=0.9, stop=['\n'], max_tokens=128):
        llm = OpenAI(model=self.llm_name, temperature=temperature, stop=['\n'], max_tokens=max_tokens)
        chain = LLMChain(llm=llm, prompt=self.prompt_temp)
        return chain.run(prompt)

def get_llm_backend(llm_name):
    if llm_name in OPENAI_CHAT_MODELS:
        return langchain_openai_chatllm(llm_name)
    elif llm_name in OPENAI_LLM_MODELS:
        return langchain_openai_llm(llm_name)
    else:
        return langchain_fastchat_llm(llm_name)


token_enc = tiktoken.get_encoding("cl100k_base")
import openai
from langchain import PromptTemplate, OpenAI, LLMChain
from transformers import AutoTokenizer, AutoModelForCausalLM
from langchain.chat_models import ChatOpenAI
from langchain.chains import LLMChain
from langchain import PromptTemplate
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from web_run.config import OPENAI_API_KEY

OPENAI_CHAT_MODELS = ["gpt-3.5-turbo","gpt-3.5-turbo-16k-0613","gpt-3.5-turbo-16k","gpt-4-0613","gpt-4-32k-0613"]
OPENAI_LLM_MODELS = ["text-davinci-003","text-ada-001"]
FASTCHAT_MODELS = ["vicuna-7b"]

class langchain_openai_chatllm:
    def __init__(self, llm_name):
        openai.api_key = OPENAI_API_KEY
        self.llm_name = llm_name
        human_template="{prompt}"
        human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
        self.chat_prompt = ChatPromptTemplate.from_messages([human_message_prompt])

    def run(self, prompt, temperature=1, stop=['\n'], max_tokens=128):
        chat = ChatOpenAI(model=self.llm_name, temperature=temperature, stop=stop, max_tokens=max_tokens)
        self.chain = LLMChain(llm=chat, prompt=self.chat_prompt)
        return self.chain.run(prompt)

class langchain_openai_llm:
    def __init__(self, llm_name):
        openai.api_key = OPENAI_API_KEY
        self.prompt_temp = PromptTemplate(
            input_variables=["prompt"], template="{prompt}"
        )
        self.llm_name = llm_name

    def run(self, prompt, temperature=0.9, stop=['\n'], max_tokens=128):
        llm = OpenAI(model=self.llm_name, temperature=temperature, stop=['\n'], max_tokens=max_tokens)
        chain = LLMChain(llm=llm, prompt=self.prompt_temp)
        return chain.run(prompt)

class langchain_fastchat_llm:
    def __init__(self, llm_name):
        openai.api_key = "EMPTY" # Not support yet
        openai.api_base = "http://localhost:8000/v1"
        self.prompt_temp = PromptTemplate(


from web_run.config import OPENAI_API_KEY

OPENAI_CHAT_MODELS = ["gpt-3.5-turbo","gpt-3.5-turbo-16k-0613","gpt-3.5-turbo-16k","gpt-4-0613","gpt-4-32k-0613"]
OPENAI_LLM_MODELS = ["text-davinci-003","text-ada-001"]
FASTCHAT_MODELS = ["vicuna-7b"]

class langchain_openai_chatllm:
    def __init__(self, llm_name):
        openai.api_key = OPENAI_API_KEY
        self.llm_name = llm_name
        human_template="{prompt}"
        human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
        self.chat_prompt = ChatPromptTemplate.from_messages([human_message_prompt])

    def run(self, prompt, temperature=1, stop=['\n'], max_tokens=128):
        chat = ChatOpenAI(model=self.llm_name, temperature=temperature, stop=stop, max_tokens=max_tokens)
        self.chain = LLMChain(llm=chat, prompt=self.chat_prompt)
        return self.chain.run(prompt)

class langchain_openai_llm:
    def __init__(self, llm_name):
        openai.api_key = OPENAI_API_KEY
        self.prompt_temp = PromptTemplate(
            input_variables=["prompt"], template="{prompt}"
        )
        self.llm_name = llm_name

    def run(self, prompt, temperature=0.9, stop=['\n'], max_tokens=128):
        llm = OpenAI(model=self.llm_name, temperature=temperature, stop=['\n'], max_tokens=max_tokens)
        chain = LLMChain(llm=llm, prompt=self.prompt_temp)
        return chain.run(prompt)

class langchain_fastchat_llm:
    def __init__(self, llm_name):
        openai.api_key = "EMPTY" # Not support yet
        openai.api_base = "http://localhost:8000/v1"
        self.prompt_temp = PromptTemplate(
            input_variables=["prompt"], template="{prompt}"
        )
        self.llm_name = llm_name

    def run(self, prompt, temperature=0.9, stop=['\n'], max_tokens=128):
        llm = OpenAI(model=self.llm_name, temperature=temperature, stop=['\n'], max_tokens=max_tokens)
        chain = LLMChain(llm=llm, prompt=self.prompt_temp)
        return chain.run(prompt)

def get_llm_backend(llm_name):
    if llm_name in OPENAI_CHAT_MODELS:
        return langchain_openai_chatllm(llm_name)
    elif llm_name in OPENAI_LLM_MODELS:


    def run(self, prompt, temperature=1, stop=['\n'], max_tokens=128):
        chat = ChatOpenAI(model=self.llm_name, temperature=temperature, stop=stop, max_tokens=max_tokens)
        self.chain = LLMChain(llm=chat, prompt=self.chat_prompt)
        return self.chain.run(prompt)

class langchain_openai_llm:
    def __init__(self, llm_name):
        openai.api_key = OPENAI_API_KEY
        self.prompt_temp = PromptTemplate(
            input_variables=["prompt"], template="{prompt}"
        )
        self.llm_name = llm_name

    def run(self, prompt, temperature=0.9, stop=['\n'], max_tokens=128):
        llm = OpenAI(model=self.llm_name, temperature=temperature, stop=['\n'], max_tokens=max_tokens)
        chain = LLMChain(llm=llm, prompt=self.prompt_temp)
        return chain.run(prompt)

class langchain_fastchat_llm:
    def __init__(self, llm_name):
        openai.api_key = "EMPTY" # Not support yet
        openai.api_base = "http://localhost:8000/v1"
        self.prompt_temp = PromptTemplate(
            input_variables=["prompt"], template="{prompt}"
        )
        self.llm_name = llm_name

    def run(self, prompt, temperature=0.9, stop=['\n'], max_tokens=128):
        llm = OpenAI(model=self.llm_name, temperature=temperature, stop=['\n'], max_tokens=max_tokens)
        chain = LLMChain(llm=llm, prompt=self.prompt_temp)
        return chain.run(prompt)

def get_llm_backend(llm_name):
    if llm_name in OPENAI_CHAT_MODELS:
        return langchain_openai_chatllm(llm_name)
    elif llm_name in OPENAI_LLM_MODELS:
        return langchain_openai_llm(llm_name)
    else:
        return langchain_fastchat_llm(llm_name)


    def run(self, prompt, temperature=0.9, stop=['\n'], max_tokens=128):
        llm = OpenAI(model=self.llm_name, temperature=temperature, stop=['\n'], max_tokens=max_tokens)
        chain = LLMChain(llm=llm, prompt=self.prompt_temp)
        return chain.run(prompt)

class langchain_fastchat_llm:
    def __init__(self, llm_name):
        openai.api_key = "EMPTY" # Not support yet
        openai.api_base = "http://localhost:8000/v1"
        self.prompt_temp = PromptTemplate(
            input_variables=["prompt"], template="{prompt}"
        )
        self.llm_name = llm_name

    def run(self, prompt, temperature=0.9, stop=['\n'], max_tokens=128):
        llm = OpenAI(model=self.llm_name, temperature=temperature, stop=['\n'], max_tokens=max_tokens)
        chain = LLMChain(llm=llm, prompt=self.prompt_temp)
        return chain.run(prompt)

def get_llm_backend(llm_name):
    if llm_name in OPENAI_CHAT_MODELS:
        return langchain_openai_chatllm(llm_name)
    elif llm_name in OPENAI_LLM_MODELS:
        return langchain_openai_llm(llm_name)
    else:
        return langchain_fastchat_llm(llm_name)


run

        self.prompt_temp = PromptTemplate(
            input_variables=["prompt"], template="{prompt}"
        )
        self.llm_name = llm_name

    def run(self, prompt, temperature=0.9, stop=['\n'], max_tokens=128):
        llm = OpenAI(model=self.llm_name, temperature=temperature, stop=['\n'], max_tokens=max_tokens)
        chain = LLMChain(llm=llm, prompt=self.prompt_temp)
        return chain.run(prompt)

class langchain_fastchat_llm:
    def __init__(self, llm_name):
        openai.api_key = "EMPTY" # Not support yet
        openai.api_base = "http://localhost:8000/v1"
        self.prompt_temp = PromptTemplate(
            input_variables=["prompt"], template="{prompt}"
        )
        self.llm_name = llm_name

    def run(self, prompt, temperature=0.9, stop=['\n'], max_tokens=128):
        llm = OpenAI(model=self.llm_name, temperature=temperature, stop=['\n'], max_tokens=max_tokens)
        chain = LLMChain(llm=llm, prompt=self.prompt_temp)
        return chain.run(prompt)

def get_llm_backend(llm_name):
    if llm_name in OPENAI_CHAT_MODELS:
        return langchain_openai_chatllm(llm_name)
    elif llm_name in OPENAI_LLM_MODELS:
        return langchain_openai_llm(llm_name)
    else:
        return langchain_fastchat_llm(llm_name)


webshop.test

import re
def idx_parser(text):
    pattern = r"\d\s\w*"  # Matches a word boundary, followed by a digit, and then any word characters

    match = re.search(pattern, text)
    if match:
        return int(match.group(0))
    else:
        return None

text = "5\n\nAction: Buy Now"
idx = idx_parser(text)
print(idx)


import logging
import pytest
import random
import shutil
from pathlib import Path
from web_agent_site.utils import *

def test_random_idx():
    random.seed(24)
    weights = [random.randint(0, 10) for _ in range(0, 50)]
    cml_weights = [0]
    for w in weights:
        cml_weights.append(cml_weights[-1] + w)
    idx_1, expected_1 = random_idx(cml_weights), 44
    idx_2, expected_2 = random_idx(cml_weights), 15
    idx_3, expected_3 = random_idx(cml_weights), 36
    assert idx_1 == expected_1
    assert idx_2 == expected_2
    assert idx_3 == expected_3

def test_setup_logger():
    LOG_DIR = 'user_session_logs_test/'
    user_log_dir = Path(LOG_DIR)
    user_log_dir.mkdir(parents=True, exist_ok=True)
    session_id = "ABC"

    logger = setup_logger(session_id, user_log_dir)
    log_file = Path(LOG_DIR + "/" + session_id + ".jsonl")
    assert Path(log_file).is_file()
    assert logger.level == logging.INFO

    content = "Hello there"
    logger.info(content)
    assert log_file.read_text().strip("\n") == content

    shutil.rmtree(LOG_DIR)

def test_generate_mturk_code():
    suite = [
        ('', 'DA39A3EE5E'),
        ('ABC', '3C01BDBB26'),
        ('123', '40BD001563'),
        ('1A1', '10E7DB0A44'),
        ('$%^ABC', '5D5607D24E')
    ]
    for session_id, expected in suite:
        output = generate_mturk_code(session_id)
        assert type(expected) is str
        assert output == expected


import logging
import pytest
import random
import shutil
from pathlib import Path
from web_agent_site.utils import *

def test_random_idx():
    random.seed(24)
    weights = [random.randint(0, 10) for _ in range(0, 50)]
    cml_weights = [0]
    for w in weights:
        cml_weights.append(cml_weights[-1] + w)
    idx_1, expected_1 = random_idx(cml_weights), 44
    idx_2, expected_2 = random_idx(cml_weights), 15
    idx_3, expected_3 = random_idx(cml_weights), 36
    assert idx_1 == expected_1
    assert idx_2 == expected_2
    assert idx_3 == expected_3

def test_setup_logger():
    LOG_DIR = 'user_session_logs_test/'
    user_log_dir = Path(LOG_DIR)
    user_log_dir.mkdir(parents=True, exist_ok=True)
    session_id = "ABC"

    logger = setup_logger(session_id, user_log_dir)
    log_file = Path(LOG_DIR + "/" + session_id + ".jsonl")
    assert Path(log_file).is_file()
    assert logger.level == logging.INFO

    content = "Hello there"
    logger.info(content)
    assert log_file.read_text().strip("\n") == content

    shutil.rmtree(LOG_DIR)

def test_generate_mturk_code():
    suite = [


    assert idx_2 == expected_2
    assert idx_3 == expected_3

def test_setup_logger():
    LOG_DIR = 'user_session_logs_test/'
    user_log_dir = Path(LOG_DIR)
    user_log_dir.mkdir(parents=True, exist_ok=True)
    session_id = "ABC"

    logger = setup_logger(session_id, user_log_dir)
    log_file = Path(LOG_DIR + "/" + session_id + ".jsonl")
    assert Path(log_file).is_file()
    assert logger.level == logging.INFO

    content = "Hello there"
    logger.info(content)
    assert log_file.read_text().strip("\n") == content

    shutil.rmtree(LOG_DIR)

def test_generate_mturk_code():
    suite = [
        ('', 'DA39A3EE5E'),
        ('ABC', '3C01BDBB26'),
        ('123', '40BD001563'),
        ('1A1', '10E7DB0A44'),
        ('$%^ABC', '5D5607D24E')
    ]
    for session_id, expected in suite:
        output = generate_mturk_code(session_id)
        assert type(expected) is str
        assert output == expected


import pytest
from math import isclose
from web_agent_site.engine.goal import *

def test_get_type_reward():
    # Exact Match
    goal = {
        'query': "Query 1",
        'product_category': "a › b › c",
        'name': "Name 1"
    }
    purchased = {
        'query': "Query 1",
        'product_category': "a › b › c",
        'name': "Name 1"
    }
    result = get_type_reward(purchased, goal)
    assert result['r_type'] == 1.
    assert result['query_match'] == True
    assert result['category_match'] == True
    assert result['title_score'] == 1

    # Query Mismatch
    purchased['query'] = 'Query 2'
    result = get_type_reward(purchased, goal)
    assert result['query_match'] == False

    # Out of order / non-matching / partially matching / duplicate categories
    purchased['product_category'] = "b › c › a"
    result = get_type_reward(purchased, goal)
    assert result['category_match'] == True
    purchased['product_category'] = "d › e › f"
    result = get_type_reward(purchased, goal)
    assert result['category_match'] == False
    purchased['product_category'] = "a › d › b"
    result = get_type_reward(purchased, goal)
    assert result['category_match'] == True
    purchased['product_category'] = "a › a › b"
    result = get_type_reward(purchased, goal)
    assert result['category_match'] == True
    purchased['product_category'] = "a › a › d"
    result = get_type_reward(purchased, goal)
    assert result['category_match'] == False

    # Similar product names
    goal['name'] = "Mens D.O.N. Issue 2 Gca Basketball Sneakers Shoes Casual - Off White"
    purchased['name'] = "PEAK High Top Mens Basketball Shoes Lou Williams Streetball Master Breathable Non Slip Outdoor Sneakers"
    result = get_type_reward(purchased, goal)
    assert isclose(result['title_score'], 0.333, abs_tol=1e-2)

    # Slightly similar product names
    goal['name'] = "Saireed UL Listed 2 Prong Power Cord for JBL Bar 3.1 Bar 2.1 Channel 4K Ultra HD Soundbar Home Theater System Subwoofer"
    purchased['name'] = "BRST AC Power Cord Outlet Socket Cable Plug Lead for Panasonic SC-HT830V DVD/VCR Combo Home Theater System"
    result = get_type_reward(purchased, goal)
    assert isclose(result['title_score'], 0.3, abs_tol=1e-2)

    goal['name'] = "Saireed UL Listed 2 Prong Power Cord for JBL Bar 3.1 Bar 2.1 Channel 4K Ultra HD Soundbar"
    purchased['name'] = "BRST AC Power Cord Outlet Socket Cable Plug Lead for Panasonic SC-HT830V DVD/VCR Combo Home Theater System"
    result = get_type_reward(purchased, goal)
    assert isclose(result['title_score'], 0.15, abs_tol=1e-2)

    # Completely different product names
    goal['name'] = "Rusticware 921ORB Kitchen and Bath Cabinet Knob"
    purchased['name'] = "Minkissy 2pcs Stainless Steel Eyebrow Tweezers Blackhead Acne Remover Portable Makeup Tweezers (Silver)"
    result = get_type_reward(purchased, goal)
    assert result['title_score'] < 0.05

def test_get_attribute_reward():
    # Exact Match
    goal = {
        'attributes': ["tea tree", "essential oils", "natural ingredients"],
    }
    purchased = {
        'Attributes': ["tea tree", "essential oil", "natural ingredients"]
    }
    r_attr, num_attr_matches = get_attribute_reward(purchased, goal)
    assert r_attr == 1
    assert num_attr_matches == 3

    # Partial Match
    goal = {
        'attributes': ["tea tree", "essential oils", "natural ingredients"]
    }
    purchased = {
        'Attributes': ["essential oil", "natural ingredients"],
        'Title': "",
        'BulletPoints': [],
        'Description': ""
    }
    r_attr, num_attr_matches = get_attribute_reward(purchased, goal)
    assert r_attr == 2./3.
    assert num_attr_matches == 2

    # Goal attributes found in purchased non-goals
    goal = {
        'attributes': ["tea tree", "essential oils", "natural ingredients"]
    }
    purchased = {
        'Attributes': [],
        'Title': "",
        'BulletPoints': ["This shampoo has essential oils and smells like lemons"],
        'Description': "Best shampoo on the market, made with natural ingredients"
    }
    r_attr, num_attr_matches = get_attribute_reward(purchased, goal)
    assert r_attr == 2./3.
    assert num_attr_matches == 2

    # No match
    goal = {
        'attributes': ["tea tree", "essential oils", "natural ingredients"]
    }
    purchased = {
        'Attributes': ["tea bag", "earl gray", "lipton"],
        'Title': "English tea for breakfast",
        'BulletPoints': ["Soothing aroma", "Calming, great feeling"],
        'Description': "Best tea made by Lipton, great to pair with breakfast"
    }
    r_attr, num_attr_matches = get_attribute_reward(purchased, goal)
    assert r_attr == 0
    assert num_attr_matches == 0

def test_get_option_reward():
    # Exact Match
    goal = ["grey", "XL", "pack of 12"]
    purchased = ["pack of 12", "grey", "XL"]
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == len(goal)
    assert r_option == 1

    # Partial Match
    goal = ["grey", "XL", "pack of 12"]
    purchased = ["pack of 12", "blue", "XL"]
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == len(goal) - 1
    assert r_option == 2./3.

    # Fuzzy Match
    goal = ["cool powder snow", "XL", "pack of 12"]
    purchased = ["pack of 12", "powder snow", "XL"]
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == len(goal)
    assert r_option == 1

    # Empty Goal Options
    goal = []
    purchased = ["goal 1", "goal 2"]
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == 0
    assert r_option == None

    # Empty Purchased Options
    goal = ["goal 1", "goal 2"]
    purchased = []
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == 0
    assert r_option == 0

def test_get_reward():
    # Exact Match
    goal = {
        'query': "Query 1",
        'product_category': "a › b › c",
        'name': "Mens D.O.N. Issue 2 Gca Basketball Sneakers Shoes Casual - Off White",
        'attributes': ["tea tree", "essential oils", "natural ingredients"],
        'goal_options': {"color": "grey", "size": "XL"},
        'price_upper': 40.00
    }
    purchased = {
        'query': "Query 1",
        'product_category': "a › b › c",
        'name': "Mens D.O.N. Issue 2 Gca Basketball Sneakers Shoes Casual - Off White",
        'Attributes': ["tea tree", "essential oil", "natural ingredients"],
        'Title': "",
        'BulletPoints': [],
        'Description': "",
        'goal_options': {"color": "grey", "size": "XL"}
    }
    total_reward = get_reward(purchased, goal, 35, purchased['goal_options'])
    assert total_reward == 1

    # Variation in r_attributes reward
    purchased['Attributes'] = []
    purchased['Title'] = ""
    purchased['BulletPoints'] = "This shampoo has essential oils and smells like lemons"
    purchased['Description'] = "Best shampoo on the market, made with natural ingredients"
    total_reward = get_reward(purchased, goal, 35, purchased['goal_options'])
    assert isclose(total_reward, 2./3., abs_tol=1e-2)

    # Variation in r_option reward
    goal['goal_options'] = {"color": "grey", "size": "XL", "amount": "pack of 12"}
    total_reward = get_reward(purchased, goal, 35, purchased['goal_options'])
    assert isclose(total_reward, 0.5714, abs_tol=1e-2)

    # Variation in r_type reward
    goal['name'] = "Saireed UL Listed 2 Prong Power Cord for JBL Bar 3.1 Bar 2.1 Channel 4K Ultra HD Soundbar"
    purchased['name'] = "BRST AC Power Cord Outlet Socket Cable Plug Lead for Panasonic SC-HT830V DVD/VCR Combo Home Theater System"
    purchased['query'] = "Query 2"
    purchased['product_category'] = "a › d › e"
    total_reward = get_reward(purchased, goal, 35, purchased['goal_options'])
    assert isclose(total_reward, 0.2857, abs_tol=1e-2)


import pytest
from math import isclose
from web_agent_site.engine.goal import *

def test_get_type_reward():
    # Exact Match
    goal = {
        'query': "Query 1",
        'product_category': "a › b › c",
        'name': "Name 1"
    }
    purchased = {
        'query': "Query 1",
        'product_category': "a › b › c",
        'name': "Name 1"
    }
    result = get_type_reward(purchased, goal)
    assert result['r_type'] == 1.
    assert result['query_match'] == True
    assert result['category_match'] == True
    assert result['title_score'] == 1

    # Query Mismatch
    purchased['query'] = 'Query 2'
    result = get_type_reward(purchased, goal)
    assert result['query_match'] == False

    # Out of order / non-matching / partially matching / duplicate categories
    purchased['product_category'] = "b › c › a"
    result = get_type_reward(purchased, goal)
    assert result['category_match'] == True
    purchased['product_category'] = "d › e › f"
    result = get_type_reward(purchased, goal)
    assert result['category_match'] == False
    purchased['product_category'] = "a › d › b"
    result = get_type_reward(purchased, goal)
    assert result['category_match'] == True
    purchased['product_category'] = "a › a › b"
    result = get_type_reward(purchased, goal)
    assert result['category_match'] == True
    purchased['product_category'] = "a › a › d"
    result = get_type_reward(purchased, goal)
    assert result['category_match'] == False

    # Similar product names
    goal['name'] = "Mens D.O.N. Issue 2 Gca Basketball Sneakers Shoes Casual - Off White"
    purchased['name'] = "PEAK High Top Mens Basketball Shoes Lou Williams Streetball Master Breathable Non Slip Outdoor Sneakers"
    result = get_type_reward(purchased, goal)
    assert isclose(result['title_score'], 0.333, abs_tol=1e-2)

    # Slightly similar product names
    goal['name'] = "Saireed UL Listed 2 Prong Power Cord for JBL Bar 3.1 Bar 2.1 Channel 4K Ultra HD Soundbar Home Theater System Subwoofer"
    purchased['name'] = "BRST AC Power Cord Outlet Socket Cable Plug Lead for Panasonic SC-HT830V DVD/VCR Combo Home Theater System"
    result = get_type_reward(purchased, goal)
    assert isclose(result['title_score'], 0.3, abs_tol=1e-2)

    goal['name'] = "Saireed UL Listed 2 Prong Power Cord for JBL Bar 3.1 Bar 2.1 Channel 4K Ultra HD Soundbar"
    purchased['name'] = "BRST AC Power Cord Outlet Socket Cable Plug Lead for Panasonic SC-HT830V DVD/VCR Combo Home Theater System"
    result = get_type_reward(purchased, goal)
    assert isclose(result['title_score'], 0.15, abs_tol=1e-2)

    # Completely different product names
    goal['name'] = "Rusticware 921ORB Kitchen and Bath Cabinet Knob"
    purchased['name'] = "Minkissy 2pcs Stainless Steel Eyebrow Tweezers Blackhead Acne Remover Portable Makeup Tweezers (Silver)"
    result = get_type_reward(purchased, goal)
    assert result['title_score'] < 0.05

def test_get_attribute_reward():
    # Exact Match
    goal = {
        'attributes': ["tea tree", "essential oils", "natural ingredients"],
    }
    purchased = {
        'Attributes': ["tea tree", "essential oil", "natural ingredients"]
    }
    r_attr, num_attr_matches = get_attribute_reward(purchased, goal)
    assert r_attr == 1
    assert num_attr_matches == 3

    # Partial Match
    goal = {
        'attributes': ["tea tree", "essential oils", "natural ingredients"]
    }
    purchased = {
        'Attributes': ["essential oil", "natural ingredients"],
        'Title': "",


    result = get_type_reward(purchased, goal)
    assert isclose(result['title_score'], 0.333, abs_tol=1e-2)

    # Slightly similar product names
    goal['name'] = "Saireed UL Listed 2 Prong Power Cord for JBL Bar 3.1 Bar 2.1 Channel 4K Ultra HD Soundbar Home Theater System Subwoofer"
    purchased['name'] = "BRST AC Power Cord Outlet Socket Cable Plug Lead for Panasonic SC-HT830V DVD/VCR Combo Home Theater System"
    result = get_type_reward(purchased, goal)
    assert isclose(result['title_score'], 0.3, abs_tol=1e-2)

    goal['name'] = "Saireed UL Listed 2 Prong Power Cord for JBL Bar 3.1 Bar 2.1 Channel 4K Ultra HD Soundbar"
    purchased['name'] = "BRST AC Power Cord Outlet Socket Cable Plug Lead for Panasonic SC-HT830V DVD/VCR Combo Home Theater System"
    result = get_type_reward(purchased, goal)
    assert isclose(result['title_score'], 0.15, abs_tol=1e-2)

    # Completely different product names
    goal['name'] = "Rusticware 921ORB Kitchen and Bath Cabinet Knob"
    purchased['name'] = "Minkissy 2pcs Stainless Steel Eyebrow Tweezers Blackhead Acne Remover Portable Makeup Tweezers (Silver)"
    result = get_type_reward(purchased, goal)
    assert result['title_score'] < 0.05

def test_get_attribute_reward():
    # Exact Match
    goal = {
        'attributes': ["tea tree", "essential oils", "natural ingredients"],
    }
    purchased = {
        'Attributes': ["tea tree", "essential oil", "natural ingredients"]
    }
    r_attr, num_attr_matches = get_attribute_reward(purchased, goal)
    assert r_attr == 1
    assert num_attr_matches == 3

    # Partial Match
    goal = {
        'attributes': ["tea tree", "essential oils", "natural ingredients"]
    }
    purchased = {
        'Attributes': ["essential oil", "natural ingredients"],
        'Title': "",
        'BulletPoints': [],
        'Description': ""
    }
    r_attr, num_attr_matches = get_attribute_reward(purchased, goal)
    assert r_attr == 2./3.
    assert num_attr_matches == 2

    # Goal attributes found in purchased non-goals
    goal = {
        'attributes': ["tea tree", "essential oils", "natural ingredients"]
    }
    purchased = {
        'Attributes': [],
        'Title': "",
        'BulletPoints': ["This shampoo has essential oils and smells like lemons"],
        'Description': "Best shampoo on the market, made with natural ingredients"
    }
    r_attr, num_attr_matches = get_attribute_reward(purchased, goal)
    assert r_attr == 2./3.
    assert num_attr_matches == 2

    # No match
    goal = {
        'attributes': ["tea tree", "essential oils", "natural ingredients"]
    }
    purchased = {
        'Attributes': ["tea bag", "earl gray", "lipton"],
        'Title': "English tea for breakfast",
        'BulletPoints': ["Soothing aroma", "Calming, great feeling"],
        'Description': "Best tea made by Lipton, great to pair with breakfast"
    }
    r_attr, num_attr_matches = get_attribute_reward(purchased, goal)
    assert r_attr == 0
    assert num_attr_matches == 0

def test_get_option_reward():
    # Exact Match
    goal = ["grey", "XL", "pack of 12"]
    purchased = ["pack of 12", "grey", "XL"]
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == len(goal)
    assert r_option == 1

    # Partial Match
    goal = ["grey", "XL", "pack of 12"]
    purchased = ["pack of 12", "blue", "XL"]
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == len(goal) - 1
    assert r_option == 2./3.

    # Fuzzy Match
    goal = ["cool powder snow", "XL", "pack of 12"]
    purchased = ["pack of 12", "powder snow", "XL"]
    r_option, matches = get_option_reward(purchased, goal)


        'Description': "Best shampoo on the market, made with natural ingredients"
    }
    r_attr, num_attr_matches = get_attribute_reward(purchased, goal)
    assert r_attr == 2./3.
    assert num_attr_matches == 2

    # No match
    goal = {
        'attributes': ["tea tree", "essential oils", "natural ingredients"]
    }
    purchased = {
        'Attributes': ["tea bag", "earl gray", "lipton"],
        'Title': "English tea for breakfast",
        'BulletPoints': ["Soothing aroma", "Calming, great feeling"],
        'Description': "Best tea made by Lipton, great to pair with breakfast"
    }
    r_attr, num_attr_matches = get_attribute_reward(purchased, goal)
    assert r_attr == 0
    assert num_attr_matches == 0

def test_get_option_reward():
    # Exact Match
    goal = ["grey", "XL", "pack of 12"]
    purchased = ["pack of 12", "grey", "XL"]
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == len(goal)
    assert r_option == 1

    # Partial Match
    goal = ["grey", "XL", "pack of 12"]
    purchased = ["pack of 12", "blue", "XL"]
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == len(goal) - 1
    assert r_option == 2./3.

    # Fuzzy Match
    goal = ["cool powder snow", "XL", "pack of 12"]
    purchased = ["pack of 12", "powder snow", "XL"]
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == len(goal)
    assert r_option == 1

    # Empty Goal Options
    goal = []
    purchased = ["goal 1", "goal 2"]
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == 0
    assert r_option == None

    # Empty Purchased Options
    goal = ["goal 1", "goal 2"]
    purchased = []
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == 0
    assert r_option == 0

def test_get_reward():
    # Exact Match
    goal = {
        'query': "Query 1",
        'product_category': "a › b › c",
        'name': "Mens D.O.N. Issue 2 Gca Basketball Sneakers Shoes Casual - Off White",
        'attributes': ["tea tree", "essential oils", "natural ingredients"],
        'goal_options': {"color": "grey", "size": "XL"},
        'price_upper': 40.00
    }
    purchased = {
        'query': "Query 1",
        'product_category': "a › b › c",
        'name': "Mens D.O.N. Issue 2 Gca Basketball Sneakers Shoes Casual - Off White",
        'Attributes': ["tea tree", "essential oil", "natural ingredients"],
        'Title': "",
        'BulletPoints': [],
        'Description': "",
        'goal_options': {"color": "grey", "size": "XL"}


    goal = ["cool powder snow", "XL", "pack of 12"]
    purchased = ["pack of 12", "powder snow", "XL"]
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == len(goal)
    assert r_option == 1

    # Empty Goal Options
    goal = []
    purchased = ["goal 1", "goal 2"]
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == 0
    assert r_option == None

    # Empty Purchased Options
    goal = ["goal 1", "goal 2"]
    purchased = []
    r_option, matches = get_option_reward(purchased, goal)
    assert matches == 0
    assert r_option == 0

def test_get_reward():
    # Exact Match
    goal = {
        'query': "Query 1",
        'product_category': "a › b › c",
        'name': "Mens D.O.N. Issue 2 Gca Basketball Sneakers Shoes Casual - Off White",
        'attributes': ["tea tree", "essential oils", "natural ingredients"],
        'goal_options': {"color": "grey", "size": "XL"},
        'price_upper': 40.00
    }
    purchased = {
        'query': "Query 1",
        'product_category': "a › b › c",
        'name': "Mens D.O.N. Issue 2 Gca Basketball Sneakers Shoes Casual - Off White",
        'Attributes': ["tea tree", "essential oil", "natural ingredients"],
        'Title': "",
        'BulletPoints': [],
        'Description': "",
        'goal_options': {"color": "grey", "size": "XL"}
    }
    total_reward = get_reward(purchased, goal, 35, purchased['goal_options'])
    assert total_reward == 1

    # Variation in r_attributes reward
    purchased['Attributes'] = []
    purchased['Title'] = ""
    purchased['BulletPoints'] = "This shampoo has essential oils and smells like lemons"
    purchased['Description'] = "Best shampoo on the market, made with natural ingredients"
    total_reward = get_reward(purchased, goal, 35, purchased['goal_options'])
    assert isclose(total_reward, 2./3., abs_tol=1e-2)

    # Variation in r_option reward
    goal['goal_options'] = {"color": "grey", "size": "XL", "amount": "pack of 12"}
    total_reward = get_reward(purchased, goal, 35, purchased['goal_options'])
    assert isclose(total_reward, 0.5714, abs_tol=1e-2)

    # Variation in r_type reward
    goal['name'] = "Saireed UL Listed 2 Prong Power Cord for JBL Bar 3.1 Bar 2.1 Channel 4K Ultra HD Soundbar"
    purchased['name'] = "BRST AC Power Cord Outlet Socket Cable Plug Lead for Panasonic SC-HT830V DVD/VCR Combo Home Theater System"
    purchased['query'] = "Query 2"
    purchased['product_category'] = "a › d › e"
    total_reward = get_reward(purchased, goal, 35, purchased['goal_options'])
    assert isclose(total_reward, 0.2857, abs_tol=1e-2)


import pytest
from web_agent_site.engine.normalize import *

def test_normalize_color():
    suite = [
        ("", ""),
        ("black forest", "black"),
        ("violet lavender", "lavender"),
        ("steelivy fuchsia", "fuchsia"),
        ("123alabaster", "alabaster"),
        ("webshop", "webshop")
    ]
    for color_string, expected in suite:
        output = normalize_color(color_string)
        assert type(output) is str
        assert output == expected

def test_normalize_color_size():
    product_prices = {
        (1, "black forest", "3 meter"): 10.29,
        (2, "violet lavender", "xx-large"): 23.42,
        (3, "steelivy fuchsia", "random value"): 193.87,
        (4, "123alabaster", "40cm plus"): 67.23,
        (5, "webshop", "142"): 1.02,
        (6, "webshopsteel", "2 petite"): 57.99,
        (7, "leather black", "91ft walnut feet"): 6.20,
    }
    color_mapping_expected = {
        'N.A.': 'not_matched',
        "black forest": "black",
        "violet lavender": "lavender",
        "steelivy fuchsia": "fuchsia",
        "123alabaster": "alabaster",
        "webshop": "not_matched",
        "webshopsteel": "steel",
        "leather black": "black"
    }
    size_mapping_expected = {
        'N.A.': 'not_matched',
        "3 meter": '(.*)meter',
        "xx-large": 'xx-large',
        "random value": "not_matched",
        "40cm plus": '(.*)plus',
        "142": "numeric_size",
        "2 petite": "(.*)petite",
        "91ft walnut feet": '(.*)ft',
    }

    color_mapping, size_mapping = normalize_color_size(product_prices)
    assert type(color_mapping) == dict
    assert type(size_mapping)  == dict
    assert color_mapping == color_mapping_expected
    assert size_mapping  == size_mapping_expected


import pytest
from web_agent_site.engine.normalize import *

def test_normalize_color():
    suite = [
        ("", ""),
        ("black forest", "black"),
        ("violet lavender", "lavender"),
        ("steelivy fuchsia", "fuchsia"),
        ("123alabaster", "alabaster"),
        ("webshop", "webshop")
    ]
    for color_string, expected in suite:
        output = normalize_color(color_string)
        assert type(output) is str
        assert output == expected

def test_normalize_color_size():
    product_prices = {
        (1, "black forest", "3 meter"): 10.29,
        (2, "violet lavender", "xx-large"): 23.42,
        (3, "steelivy fuchsia", "random value"): 193.87,
        (4, "123alabaster", "40cm plus"): 67.23,
        (5, "webshop", "142"): 1.02,
        (6, "webshopsteel", "2 petite"): 57.99,
        (7, "leather black", "91ft walnut feet"): 6.20,
    }
    color_mapping_expected = {
        'N.A.': 'not_matched',
        "black forest": "black",
        "violet lavender": "lavender",
        "steelivy fuchsia": "fuchsia",
        "123alabaster": "alabaster",
        "webshop": "not_matched",
        "webshopsteel": "steel",
        "leather black": "black"


import json
import pytest
import requests_mock

from transfer.predict_help import (
    convert_dict_to_actions,
    parse_item_page_amz, parse_results_amz,
    parse_item_page_ebay, parse_results_ebay,
    parse_item_page_ws, parse_results_ws,
    Page, WEBSHOP_URL, WEBSHOP_SESSION
)

@requests_mock.Mocker(kw="mock")
def test_parse_item_page_ws(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_item_page_ws", "rb")
    mock_body = mock_file.read()
    mock_file.close()

    mock_desc_file = open("tests/transfer/mocks/mock_parse_item_page_ws_desc", "rb")
    mock_desc_body = mock_desc_file.read()
    mock_desc_file.close()

    mock_feat_file = open("tests/transfer/mocks/mock_parse_item_page_ws_feat", "rb")
    mock_feat_body = mock_feat_file.read()
    mock_feat_file.close()

    mock_asin = "B09P87V3LZ"
    mock_query = "red basketball shoes"
    mock_options = {}

    # Invoke function, check response
    query_str = '+'.join(mock_query.split())
    options_str = json.dumps(mock_options)
    url = (
        f"{WEBSHOP_URL}/item_page/{WEBSHOP_SESSION}/"
        f"{mock_asin}/{query_str}/1/{options_str}"
    )
    url_desc = (
        f"{WEBSHOP_URL}/item_sub_page/{WEBSHOP_SESSION}/"
        f"{mock_asin}/{query_str}/1/Description/{options_str}"
    )
    url_feat = (
        f"{WEBSHOP_URL}/item_sub_page/{WEBSHOP_SESSION}/"
        f"{mock_asin}/{query_str}/1/Features/{options_str}"
    )
    print(f"Item Page URL: {url}")
    print(f"Item Description URL: {url_desc}")
    print(f"Item Features URL: {url_feat}")

    kwargs["mock"].get(url, content=mock_body)
    kwargs["mock"].get(url_desc, content=mock_desc_body)
    kwargs["mock"].get(url_feat, content=mock_feat_body)

    output = parse_item_page_ws(mock_asin, mock_query, 1, mock_options)
    expected = {
        'MainImage': 'https://m.media-amazon.com/images/I/51ltvkzGhGL.jpg',
        'Price': '100.0',
        'Rating': 'N.A.',
        'Title': 'PMUYBHF Womens Fashion Flat Shoes Comfortable Running Shoes '
            'Sneakers Tennis Athletic Shoe Casual Walking Shoes',
        'asin': mock_asin,
        'option_to_image': {
            '6.5': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27size%27:%20%276.5%27%7D',
            '7.5': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27size%27:%20%277.5%27%7D',
            '8': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27size%27:%20%278%27%7D',
            '8.5': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27size%27:%20%278.5%27%7D',
            '9': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27size%27:%20%279%27%7D',
            'black': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27color%27:%20%27black%27%7D',
            'purple': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27color%27:%20%27purple%27%7D',
            'red': 'http://3.83.245.205:3000/item_page/abc/B09P87V3LZ/%5B%27red%27%2C%20%27basketball%27%2C%20%27shoes%27%5D/1/%7B%27color%27:%20%27red%27%7D'
        },
        'options': {
            'color': ['black', 'purple', 'red'],
            'size': ['6.5', '7.5', '8', '8.5', '9']
        },
        'BulletPoints': 'Pure Running Shoe\nComfort Flat Sneakers\n[FEATURES]: Soles with unique non-slip pattern, it has great '
                'abrasion resistant and provide protection when you walking or running. (Pure Running Shoe Mesh Walking Shoes Fashion '
                'Sneakers Slip On Sneakers Wedge Platform Loafers Modern Walking Shoes Sock Sneakers Platform Loafers Shoes Non Slip '
                'Running Shoes Athletic Tennis Shoes Blade Type Sneakers Lace-up Sneaker) sole\n[WIDE ANKLE DESIGN]: Perfect accord with human body '
                'engineering, green, healthy concept design make the walking shoes wear more comfortable, wide width wlking shoes. (Low '
                'Top Walking Shoes Fashion Canvas Sneakers Slip On Shoes Casual Walking Shoes Hidden Wedge Sneaker Low Top Canvas '
                'Sneakers Lace-up Classic Casual Shoes Walking Tennis Shoes Lightweight Casual Sneakers Slip on Sock Sneakers Air '
                'Cushion Platform Loafers Slip-On Mule Sneaker )\n[CUSHION WITH ARCH SUPPORT]: Gives you a comfort for all day '
                'long. Wear these lightweight walking shoes, let every step of moving on a comfortable feeling. (Fashion Casual Shoes '
                'Athletic Workout Shoes Fitness Sneaker Athletic Running Shoes Air Cushion Sneakers Stylish Athletic Shoes Lace Up '
                'Canvas Shoes Slip on Walking Shoe Fashion Sneakers Low Top Classic Sneakers Comfort Fall Shoes Memory Foam Slip On '
                'Sneakers Air Cushion Sneakers Running Walking Shoes)\n[NON-SLIP SOLE]: Made from ultra soft and lightweight RUBBER '
                'material,with the function of shock absorbing and cushioning,offering the best durability and traction. (Wedge '
                'Sneakers Walking Tennis Shoes Slip On Running Shoes Lightweight Fashion Sneakers Fashion Travel Shoes Walking '
                'Running Shoes Non Slip Running Shoes Athletic Tennis Sneakers Sports Walking Shoes Platform Fashion Sneaker '
                'Memory Foam Tennis Sneakers Running Jogging Shoes Sock Sneakers Canvas Fashion Sneakers)\n[OCCASIONS]: Ultra lightweight design provides actual '
                'feelings of being barefooted and like walking on the feather, perfect for walking, hiking, bike riding, working, '
                'shopping, indoor, outdoor, casual, sports, travel, exercise, vacation, and etc. (Flat Fashion Sneakers Lightweight '
                'Walking Sneakers Platform Loafers Sport Running Shoes Casual Flat Loafers Slip-On Sneaker Casual Walking Shoes High Top '
                'Canvas Sneakers Lace up Sneakers Workout Walking Shoes Tennis Fitness Sneaker)\n'
                '[Customers Are Our Priority]: We follow the principle of customer first, so if you encounter any problems after '
                'buying shoes, we will try our best to solve them for you. (Breathable Air Cushion Sneakers Walking Tennis Shoes Air '
                'Athletic Running Shoes Air Cushion Shoes Mesh Sneakers Fashion Tennis Shoes Jogging Walking Sneakers Breathable '
                'Casual Sneakers Fashion Walking Shoes Athletic Running Sneakers Walking Work Shoes Air Running Shoes Slip on '
                'Sneakers Mesh Walking Shoes)',
        'Description': 'Here Are The Things You Want To Knowa─=≡Σ(((つ̀ώ)つSTORE INTRODUCTION:>>>>Our store helps our customers improve their '
               'quality of life~As a distributor, we value quality and service.Focus on the high quality and durability of the '
               'product.Committed to creating a store that satisfies and reassures our customers.TIPS:>>>>1. Please allow minor errors '
               'in the data due to manual measurements.2. Due to the color settings of the display, the actual color may be slightly '
               'different from the online image.QUALITY PROMISE:>>>>Our goal is to continuously provide a range of quality products.We '
               'place a huge emphasis on the values of quality and reliability.We have always insisted on fulfilling this '
               'commitment.In short, we want our customers to have the same great product experience every time and be trusted to deliver '
               'on this commitment.Please give us a chance to serve you.OTHER:>>>>athletic sneaker laces athletic sneakers white '
               'athletic sneakers for women clearance leather Sneaker leather sneakers women leather sneakers for menleather sneaker laces '
               'leather sneaker platform basketball shoes basketball shoes for men basketball shoe laces basketball shoe grip basketball '
               'shoes for women fitness shoes for men fitness shoes women workout fitness shoes women fitness shoes women size 5 '
               'fitness shoes men workout fitness shoes for men high top sneakers for women walking shoes sneakers with arch support for women'
    }
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_item_page_ebay(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_item_page_ebay", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_asin = "403760625150"

    # Invoke function, check response
    kwargs["mock"].get(f"https://www.ebay.com/itm/{mock_asin}", content=mock_body)
    output = parse_item_page_ebay(mock_asin)
    expected = {
        'BulletPoints': 'Item specifics Condition:New without box: A brand-new, '
            'unused, and unworn item (including handmade items) that is '
            'not in ...  Read moreabout the conditionNew without box: A '
            'brand-new, unused, and unworn item (including handmade '
            'items) that is not in original packaging or may be missing '
            'original packaging materials (such as the original box or '
            'bag). The original tags may not be attached. For example, '
            'new shoes (with absolutely no signs of wear) that are no '
            'longer in their original box fall into this category.  See '
            'all condition definitionsopens in a new window or tab  '
            'Closure:Lace Up US Shoe Size:10 Occasion:Activewear, Casual '
            'Silhouette:Puma Fabric Type:Mesh Vintage:No Cushioning '
            'Level:Moderate Department:Men Style:Sneaker Outsole '
            'Material:Rubber Features:Breathable, Comfort, Cushioned, '
            'Performance Season:Fall, Spring, Summer, Winter '
            'Idset_Mpn:193990-21 Shoe Shaft Style:Low Top Style '
            'Code:193990-16 Pattern:Solid Character:J. Cole Lining '
            'Material:Synthetic Color:Red Brand:PUMA Type:Athletic '
            'Customized:No Model:RS-Dreamer Theme:Sports Shoe '
            'Width:Standard Upper Material:Textile Insole '
            'Material:Synthetic Performance/Activity:Basketball Product '
            'Line:Puma Dreamer',
        'Description': 'N/A',
        'MainImage': 'https://i.ebayimg.com/images/g/4ggAAOSwpk1ioTWz/s-l500.jpg',
        'Price': 'N/A',
        'Rating': None,
        'Title': "Puma RS-Dreamer J. Cole Basketball Shoes Red 193990-16 Men's Size 10.0",
        'asin': '403760625150',
        'option_to_image': {},
        'options': {},
    }
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_item_page_amz(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_item_page_amz", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_asin = "B073WRF565"

    # Invoke function, check response
    kwargs["mock"].get(f"https://www.amazon.com/dp/{mock_asin}", content=mock_body)
    output = parse_item_page_amz(mock_asin)
    expected = {
        'asin': 'B073WRF565',
        'Title': 'Amazon Basics Foldable 14" Black Metal Platform Bed Frame with Tool-Free Assembly No Box Spring Needed - Full',
        'Price': 'N/A',
        'Rating': '4.8 out of 5 stars',
        'BulletPoints': ' \n About this item    '
            'Product dimensions: 75" L x 54" W x 14" H | Weight: 41.4 pounds    '
            'Designed for sleepers up to 250 pounds    Full size platform bed frame offers a quiet, noise-free, '
            'supportive foundation for a mattress. No box spring needed    Folding mechanism makes the frame easy '
            'to store and move in tight spaces    Provides extra under-the-bed storage space with a vertical clea'
            'rance of about 13 inches    \n › See more product details ',
        'Description': 'Amazon Basics Foldable, 14" Black Metal Platform Bed Frame with Tool-Free Assembly, No Box Spring Needed - Full   Amazon Basics',
        'MainImage': 'https://images-na.ssl-images-amazon.com/images/I/41WIGwt-asL.__AC_SY300_SX300_QL70_FMwebp_.jpg',
        'options': {'size': ['Twin', 'Full', 'Queen', 'King'],
        'style': ['14-Inch', '18-Inch']},
        'option_to_image': {}
    }
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_results_ebay(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_results_ebay", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_query = "red basketball shoes"

    # Invoke function, check response
    query = mock_query.replace(" ", "+")
    kwargs["mock"].get(f'https://www.ebay.com/sch/i.html?_nkw={query}&_pgn=1', content=mock_body)
    output = parse_results_ebay(mock_query, 1)
    expected = [{
		'Price': ['100.00', '150.00'],
		'Title': "Reebok Answer IV Men's Basketball Shoes",
		'asin': '175065123030'
	}, {
		'Price': '$119.90',
		'Title': "Air Jordan Stay Loyal Shoes Black Red White DB2884-001 Men's Multi "
		'Size NEW',
		'asin': '265672133690'
	}, {
		'Price': '$100.00',
		'Title': "Fila Men's Stackhouse Spaghetti Basketball Shoes Black Red White "
		'1BM01788-113',
		'asin': '175282509234'
	}, {
		'Price': ['61.99',
			'85.99'
		],
		'Title': 'Puma Disc Rebirth 19481203 Mens Black Red Synthetic Athletic '
		'Basketball Shoes',
		'asin': '313944854658'
	}, {
		'Price': '$0.01',
		'Title': "Puma RS-Dreamer J. Cole Basketball Shoes Red 193990-16 Men's Size "
		'10.0',
		'asin': '403760625150'
	}, {
		'Price': '$45.00',
		'Title': 'Nike Mens 9.5 PG 5  Maroon Red White Basketball Shoes Sneaker DM '
		'5045–601￼ Flaw',
		'asin': '115456853186'
	}, {
		'Price': ['114.90',
			'119.90'
		],
		'Title': "Air Jordan Stay Loyal Shoes White Black Red DB2884-106 Men's Multi "
		'Size NEW',
		'asin': '155046831159'
	}, {
		'Price': '$8.99',
		'Title': "Harden Volume 3 Men's Basketball Shoes Size 9.5",
		'asin': '175342407862'
	}, {
		'Price': '$59.97',
		'Title': "Men's Nike Precision 5 Basketball Shoes Gym Red Black Grey Bred "
		'Multi Size NEW',
		'asin': '134149634710'
	}]
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_results_amz(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_results_amz", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_query = "red basketball shoes"

    # Invoke function, check response
    query = mock_query.replace(" ", "+")
    kwargs["mock"].get(f"https://www.amazon.com/s?k={query}&page=1", content=mock_body)
    output = parse_results_amz(mock_query, 1)
    expected = [{
		'Price': '59.49',
		'Title': 'High Top Mens Basketball Shoes Lou Williams Streetball Master '
			'Breathable Non Slip Outdoor Sneakers Cushioning Workout Shoes for '
			'Fitness',
		'asin': 'B083QCWF61'
	}, {
		'Price': '45.99',
		'Title': 'Kids Basketball Shoes High-top Sports Shoes Sneakers Durable '
		'Lace-up Non-Slip Running Shoes Secure for Little Kids Big Kids and '
		'Boys Girls',
		'asin': 'B08FWWWQ11'
	}, {
		'Price': '64.99',
		'Title': 'Unisex-Adult Lockdown 5 Basketball Shoe',
		'asin': 'B0817BFNC4'
	}, {
		'Price': '63.75',
		'Title': 'Unisex-Child Team Hustle D 9 (Gs) Sneaker',
		'asin': 'B07HHTS79M'
	}, {
		'Price': '74.64',
		'Title': 'Unisex-Adult D.O.N. Issue 3 Basketball Shoe',
		'asin': 'B08N8DQLS2'
	}, {
		'Price': '104.90',
		'Title': "Men's Lebron Witness IV Basketball Shoes",
		'asin': 'B07TKMMHVB'
	}, {
		'Price': '36.68',
		'Title': "Unisex-Child Pre-School Jet '21 Basketball Shoe",
		'asin': 'B08N6VRHV4'
	}, {
		'Price': '59.98',
		'Title': "Men's Triple Basketball Shoe",
		'asin': 'B08QCL8VKM'
	}, {
		'Price': '45.98',
		'Title': 'Unisex-Child Pre School Lockdown 4 Basketball Shoe',
		'asin': 'B07HKP12DH'
	}, {
		'Price': '143.72',
		'Title': "Men's Basketball Shoes",
		'asin': 'B07SNR7HRF'
	}]
    assert output == expected

@requests_mock.Mocker(kw="mock")
def test_parse_results_ws(**kwargs):
    # Read mock response data
    mock_file = open("tests/transfer/mocks/mock_parse_results_ws", "rb")
    mock_body = mock_file.read()
    mock_file.close()
    mock_query = "red basketball shoes"

    # Invoke function, check response
    query_str = mock_query.replace(" ", "+")
    url = (
        f'{WEBSHOP_URL}/search_results/{WEBSHOP_SESSION}/'
        f'{query_str}/1'
    )
    kwargs["mock"].get(url, content=mock_body)
    output = parse_results_ws(mock_query, 1)
    expected = [{
        'Price': [24.49, 39.99],
        'Title': "BinGoDug Men's Basketball Shoes, Men's Fashion Sneakers, Air "
        'Basketball Shoes for Men, Womens Basketball Shoes, Mens Basketball '
        'Shoes, Boys Basketball Shoes, Youth Basketball Shoes Men Women',
        'asin': 'B09GKFNQWT'
    }, {
        'Price': [1.89, 7.58],
        'Title': "RQWEIN Comfortable Mesh Sneakers Men's Roading Running Shoes "
        'Tennis Shoes Casual Fashion Sneakers Outdoor Non Slip Gym Athletic '
        'Sport Shoes',
        'asin': 'B09BFY2R3R'
    }, {
        'Price': 100.0,
        'Title': 'PMUYBHF Womens Fashion Flat Shoes Comfortable Running Shoes '
        'Sneakers Tennis Athletic Shoe Casual Walking Shoes',
        'asin': 'B09P87V3LZ'
    }, {
        'Price': 100.0,
        'Title': 'PMUYBHF Fashion Travel Shoes Jogging Walking Sneakers Air Cushion '
        'Platform Loafers Air Cushion Mesh Shoes Walking Dance Shoes',
        'asin': 'B09N6SNKC1'
    }, {
        'Price': 100.0,
        'Title': "PMUYBHF Women's Ballet Flats Walking Flats Shoes Dressy Work Low "
        'Wedge Arch Suport Flats Shoes Slip On Dress Shoes',
        'asin': 'B09N6X5S74'
    }, {
        'Price': 100.0,
        'Title': "PWKSELW High-top Men's Basketball Shoes Outdoor Sports Shoes "
        'Cushioning Training Shoes Casual Running Shoes',
        'asin': 'B09MDB9V5W'
    }, {
        'Price': 100.0,
        'Title': "Women's Flat Shoes Classic Round Toe Slip Office Black Ballet "
        'Flats Walking Flats Shoes Casual Ballet Flats',
        'asin': 'B09N6PDFRF'
    }, {
        'Price': 100.0,
        'Title': "Women's Mid-Calf Boots Wide Calf Boots for Women Fashion Zipper "
        'Womens Shoes Pu Leather Casual Boots Womens Slip-On Womens Flat '
        "Shoes Med Heel Womens' Boots Winter Snow Boot Comfy Boots(,5.5)",
        'asin': 'B09N8ZHFNM'
    }, {
        'Price': 100.0,
        'Title': 'PMUYBHF Womens Leisure Fitness Running Sport Warm Sneakers Shoes '
        'Slip-On Mule Sneakers Womens Mules',
        'asin': 'B09P87DWGR'
    }, {
        'Price': 100.0,
        'Title': 'Men Dress Shoes Leather Modern Classic Business Shoes Lace Up '
        'Classic Office Shoes Business Formal Shoes for Men',
        'asin': 'B09R9MMTKR'
    }]
    assert output == expected

def test_convert_dict_to_actions():
    # Test RESULTS page type
    asin = "334490012932"
    page_num = 2
    products = [{
        'asin': '125331076844',
        'Title': 'Modern Tall Torchiere Floor Lamp Brushed Nickel Chrome Metal Decor Living Room',
        'Price': '$129.95'
    }, {
        'asin': '125109985453',
        'Title': 'Floor Lamps Set of 2 Polished Steel Crystal Glass for Living Room Bedroom',
        'Price': '$179.99'
    }, {
        'asin': '125265434055',
        'Title': 'Floor Lamp Nickel/Polished Concrete Finish with Off-White Linen Fabric Shade',
        'Price': '$130.68'
    }, {
        'asin': '195197281169',
        'Title': 'New ListingVintage Mid Century Modern Glass Amber Globe Tension Pole Floor Lamp Light',
        'Price': '$165.00'
    }, {
        'asin': '195197512929',
        'Title': 'New ListingVTG Brass Floor Lamp Glass Shade 63.5" Tall 12" Diameter Gla
# ... truncated ...