Official Repo for paper MINT: Multi-turn Interactive Evaluation for Large Language Models by Xingyao Wang*, Zihan Wang*, Jiateng Liu, Yangyi Chen, Lifan Yuan, Hao Peng and Heng Ji.
import os

# from ReACT
path = "data/processed/alfworld/prompts/alfworld_3prompts.json"
output_dir = os.path.dirname(path)

import json

with open(path, "r") as f:
    data = json.load(f)

# altogether, we have 8 types
new_data = {}
for key in data.keys():
    if not key.startswith("react"):
        # skip act only
        continue

    output_filepath = os.path.join(output_dir, key + ".txt")

    current_prompt = ""
    old_prompt = data[key]
    split = old_prompt.split("\n")
    for k in range(len(split)):
        if split[k][0:6] == "Action":
            Action = split[k][7:].strip()
            if Action[0:3] == "go ":  # goto
                object = split[k].split("go to")[1].strip()
                current_prompt = (
                    current_prompt
                    + "<execute> "
                    + 'goto("'
                    + object
                    + '")'
                    + " </execute>"
                    + "\n"
                )
            elif Action[0:3] == "put":  # put
                object_1 = split[k].split("put")[1].strip().split("in/on")[0].strip()
                object_2 = split[k].split("in/on")[1].strip()
                current_prompt = (
                    current_prompt
                    + "<execute> "
                    + 'put("'
                    + object_1
                    + '","'
                    + object_2
                    + '")'
                    + " </execute>"
                    + "\n"
                )
            elif Action[0:3] == "tak":  # take
                object_1 = split[k].split("take")[1].strip().split("from")[0].strip()
                object_2 = split[k].split("from")[1].strip()
                current_prompt = (
                    current_prompt
                    + "<execute> "
                    + 'take_from("'
                    + object_1
                    + '","'
                    + object_2
                    + '")'
                    + " </execute>"
                    + "\n"
                )
            elif Action[0:3] == "ope":  # open
                object = split[k].split("open")[1].strip()
                current_prompt = (
                    current_prompt
                    + "<execute> "
                    + 'open("'
                    + object
                    + '")'
                    + " </execute>"
                    + "\n"
                )
            elif Action[0:3] == "tog":  # toggle
                object = split[k].split("toggle")[1].strip()
                current_prompt = (
                    current_prompt
                    + "<execute> "
                    + 'toggle("'
                    + object
                    + '")'
                    + " </execute>"
                    + "\n"
                )
            elif Action[0:3] == "clo":  # close
                object = split[k].split("close")[1].strip()
                current_prompt = (
                    current_prompt
                    + "<execute> "
                    + 'close("'
                    + object
                    + '")'
                    + " </execute>"
                    + "\n"
                )
            elif Action[0:3] == "cle":  # clean
                object_1 = split[k].split("clean")[1].strip().split("with")[0].strip()
                object_2 = split[k].split("with")[1].strip()
                current_prompt = (
                    current_prompt
                    + "<execute> "
                    + 'clean("'
                    + object_1
                    + '","'
                    + object_2
                    + '")'
                    + " </execute>"
                    + "\n"
                )
            elif Action[0:3] == "hea":  # heat
                object_1 = split[k].split("heat")[1].strip().split("with")[0].strip()
                object_2 = split[k].split("with")[1].strip()
                current_prompt = (
                    current_prompt
                    + "<execute> "
                    + 'heat("'
                    + object_1
                    + '","'
                    + object_2
                    + '")'
                    + " </execute>"
                    + "\n"
                )
            elif Action[0:3] == "coo":  # heat
                object_1 = split[k].split("cool")[1].strip().split("with")[0].strip()
                object_2 = split[k].split("with")[1].strip()
                current_prompt = (
                    current_prompt
                    + "<execute> "
                    + 'cool("'
                    + object_1
                    + '","'
                    + object_2
                    + '")'
                    + " </execute>"
                    + "\n"
                )
            elif Action[0:3] == "loo":  # look
                current_prompt = (
                    current_prompt + "<execute> " + "look()" + " </execute>" + "\n"
                )
            elif Action[0:3] == "use":  # use
                object = split[k].split("use")[1].strip()
                current_prompt = (
                    current_prompt
                    + "<execute> "
                    + 'use("'
                    + object
                    + '")'
                    + " </execute>"
                    + "\n"
                )
            else:
                raise Exception(f"invalid input: {split[k]}")
        elif split[k][0:6] == "Observ":
            current_prompt = current_prompt + split[k] + "\n"
        elif split[k][0:6] == "Think:":
            current_prompt = (
                current_prompt
                + "<thought> "
                + split[k][7:].strip()
                + " </thought>"
                + "\n"
            )
        else:
            current_prompt = current_prompt + split[k] + "\n"

    with open(output_filepath, "w") as f:
        f.write(current_prompt)

data/processed/alfworld/prompts/process_format.py
xingyaoww/mint-bench
from .base import LMAgent
from .openai_lm_agent import OpenAILMAgent
from .bard_agent import BardLMAgent
from .openai_feedback_agent import OpenAIFeedbackAgent
from .claude_feedback_agent import ClaudeFeedbackAgent
from .vllm_feedback_agent import VLLMFeedbackAgent
from .vllm_agent import VLLMAgent
from .claude_agent import ClaudeLMAgent

mint/agents/__init__.py
xingyaoww/mint-bench
import os
import logging
import traceback
import backoff
from typing import List, Mapping
import google.api_core.exceptions
import google.generativeai as palm
from google.generativeai.types import ChatResponse
from google.api_core.exceptions import InvalidArgument
from .base import LMAgent
from mint.datatypes import Action

palm.configure(api_key=os.environ.get("BARD_API_KEY", None))

LOGGER = logging.getLogger("MINT")


class BardIssue(Exception):
    pass


class BardLMAgent(LMAgent):
    def __init__(self, config):
        super().__init__(config)
        assert "model_name" in config.keys()
        self.stop_words = ["\nObservation:", "\nExpert feedback:", "\nTask:", "\n---"]
        if not self.config["model_name"].startswith("models/"):
            self.config["model_name"] = f"models/{self.config['model_name']}"

    def parse_bard_messages(self, messages: List[Mapping[str, str]]):
        # Bard accepts messages as a list of strings
        if self.config.get("add_system_message", False):
            messages = self.add_system_message(messages)
            assert messages[0]["role"] == "system"
            system = messages[0]["content"]
            messages = messages[1:]
            messages = [
                {"author": m["role"], "content": m["content"]} for m in messages
            ]
            return {
                "context": system,
                "examples": None,
                "messages": messages,
            }
        else:
            messages = [
                {"author": m["role"], "content": m["content"]} for m in messages
            ]
            return {
                "context": None,
                "examples": None,
                "messages": messages,
            }

    @backoff.on_exception(
        backoff.expo,
        (
            google.api_core.exceptions.GatewayTimeout,
            google.api_core.exceptions.ServiceUnavailable,
            google.api_core.exceptions.InternalServerError,
            google.api_core.exceptions.TooManyRequests,
        ),
    )
    def call_lm(self, messages: List[str]):
        parsed_prompt = self.parse_bard_messages(messages)
        token_stats = palm.count_message_tokens(
            context=parsed_prompt["context"],
            examples=parsed_prompt["examples"],
            messages=parsed_prompt["messages"],
        )

        # Prepend the prompt with the system message
        try:
            candidate_count = self.config.get("candidate_count", 1)
            response: ChatResponse = palm.chat(
                model=self.config["model_name"],
                context=parsed_prompt["context"],
                examples=parsed_prompt["examples"],
                messages=parsed_prompt["messages"],
                temperature=self.config.get("temperature", 0.0),
                candidate_count=candidate_count,
            )

            response_str = response.last
            # Best effort to find the candidate that contains <execute> or <solution>
            if candidate_count > 1:
                for candidate in response.candidates:
                    cur_content = candidate["content"]
                    if "<execute>" in cur_content or "<solution>" in cur_content:
                        response_str = cur_content
        except IndexError:
            raise BardIssue(
                f"IndexError (Likely triggered a filter). {traceback.format_exc()}"
            )
        except InvalidArgument:
            raise BardIssue(
                f"InvalidArgument (Likely exceeded maxlen). {traceback.format_exc()}"
            )
        if hasattr(response, "filters") and len(response.filters) > 0:
            raise BardIssue(
                f"Filters triggered: {response.filters}. {traceback.format_exc()}"
            )
        if (
            (not hasattr(response, "candidates"))
            or (response.candidates is None)
            or (len(response.candidates) == 0)
        ):
            raise BardIssue(f"Empty candidates: {response}")
        # cut off based on self.stop_words
        for stop_word in self.stop_words:
            if stop_word in response_str:
                response_str = response_str[: response_str.index(stop_word)].rstrip()
        LOGGER.info(token_stats)
        return response_str, token_stats

    def act(self, state):
        messages = state.history
        try:
            lm_output, token_usage = self.call_lm(messages)
            for usage_type, count in token_usage.items():
                state.token_counter[usage_type] += count
            action = self.lm_output_to_action(lm_output)
            return action
        except BardIssue as e:
            LOGGER.error(f"BardIssue\n{e.args[0]}")
            return Action(None, False, error=f"BardIssue\n{e.args[0]}")

mint/agents/bard_agent.py
xingyaoww/mint-bench
import logging
from typing import List, Dict, Any, Mapping

LOGGER = logging.getLogger("MINT")

from mint.datatypes import Action, State


class LMAgent:
    """Base class for an agent."""

    def __init__(self, config: Mapping[str, Any]):
        self.config = config
        LOGGER.info(f"Initialized {self.__class__.__name__} with config: {config}")
        # The agent should not generate observations or expert feedback
        self.stop_words = ["\nObservation:", "\nExpert feedback:", "\nTask:", "\n---"]

    def lm_output_to_action(self, lm_output: str) -> Action:
        propose_solution = bool("<solution>" in lm_output)
        return Action(lm_output, not propose_solution)

    def act(self, state: State) -> Action:
        """
        The history should be a format like:
        [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Who won the world series in 2020?"},
            {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."},
            {"role": "user", "content": "Where was it played?"}
        ]
        """
        raise NotImplementedError

    def add_system_message(
        self, messages: List[Dict[str, str]]
    ) -> List[Dict[str, str]]:
        # Prepend the prompt with the system message
        first_msg = messages[0]
        assert first_msg["role"] == "user"
        system, examples, task = first_msg["content"].split("\n---\n")
        messages = [
            {"role": "system", "content": system},
            {"role": "user", "content": examples + "\n---\n" + task},
        ] + messages[1:]
        return messages

mint/agents/base.py
xingyaoww/mint-bench
from .base import LMAgent
import logging
import traceback
from mint.datatypes import Action
import backoff
import requests
import os
import json

LOGGER = logging.getLogger("MINT")

url = "https://api.anthropic.com/v1/complete"
headers = {
    "accept": "application/json",
    "anthropic-version": "2023-06-01",
    "content-type": "application/json",
    "x-api-key": os.environ.get("ANTHROPIC_API_KEY", None),
}


class ClaudeLMAgent(LMAgent):
    def __init__(self, config):
        super().__init__(config)
        assert "model_name" in config.keys()
        self.stop_words = [
            "Observation:",
            "Expert feedback:",
            "Task:",
            "---",
            "\n\nHuman:",
        ]

    @backoff.on_exception(
        backoff.expo,
        requests.exceptions.RequestException,
    )
    def call_lm(self, messages):
        # Prepend the prompt with the system message
        data = {
            "model": self.config["model_name"],
            "prompt": "",
            "max_tokens_to_sample": self.config.get("max_tokens", 512),
            "temperature": self.config.get("temperature", 0),
            "stop_sequences": self.stop_words,
        }
        for message in messages:
            if message["role"] == "user":
                data["prompt"] += f"\n\nHuman: {message['content']}"
            else:
                data["prompt"] += f"\n\nAssistant: {message['content']}"
        assert len(messages) % 2 == 1, "messages must be odd length"
        data["prompt"] += "\n\nAssistant:"

        response = requests.post(url, headers=headers, json=data)

        if response.status_code == 200:
            pass
        else:
            raise requests.exceptions.RequestException(
                "Request failed with status code:", response.status_code
            )

        return json.loads(response.text)["completion"], {}

    def act(self, state):
        messages = state.history
        lm_output, token_usage = self.call_lm(messages)
        for usage_type, count in token_usage.items():
            state.token_counter[usage_type] += count
        action = self.lm_output_to_action(lm_output)
        return action
        # except Exception as e:
        #     tb = traceback.format_exc()
        #     return Action(None, False, error="Unknown error")

mint/agents/claude_agent.py
xingyaoww/mint-bench
import re
import logging
from .openai_lm_agent import OpenAILMAgent
from mint.datatypes import State, Action
from mint.prompt import FeedbackPromptTemplate
import logging
from mint.datatypes import Action
import backoff
import requests
import os
import json

LOGGER = logging.getLogger("MINT")

url = "https://api.anthropic.com/v1/complete"
headers = {
    "accept": "application/json",
    "anthropic-version": "2023-06-01",
    "content-type": "application/json",
    "x-api-key": os.environ.get("ANTHROPIC_API_KEY", None),
}


class ClaudeFeedbackAgent(OpenAILMAgent):
    def __init__(self, config):
        super().__init__(config)
        # The agent should not generate Assistant msg since it should provide feedback
        self.stop_words = ["\nObservation:", "\nTask:", "\nAssistant:"]
        self.feedback_prompt = FeedbackPromptTemplate()

    def lm_output_to_action(self, lm_output, form) -> Action:
        if form == "textual":
            feedback = lm_output
        elif form == "binary":
            # Find the first sentence (as feedback).
            first_sent = re.findall(r"([^.]*\.)", lm_output)[0]
            if "GOOD" in first_sent:
                feedback = "This is GOOD."
            elif "BAD" in first_sent:
                feedback = "This is BAD."
            else:
                raise ValueError(f"Cannot find GOOD or BAD in feedback: {feedback}")
        return Action(feedback, use_tool=False)

    def act(
        self,
        state: State,
        observation: str,
        form: str,
        gt,
        task_in_context_example: str,
        tool_desc: str,
    ) -> Action:
        gt_solution = (
            (
                f"Correct solution (please DO NOT disclose the correct solution to the assistant): {str(gt).strip()}\n"
            )
            if gt
            else "Correct solution (please DO NOT disclose the correct solution to the assistant): NOT GIVEN\n"
        )
        trajectory = "---\n".join(state.history[0]["content"].split("---\n")[2:]) + "\n"
        trajectory += "\n".join([x["content"] for x in state.history[1:]])
        trajectory += "\n" + observation
        trajectory = trajectory[
            trajectory.find("Task:") :
        ]  # Get rid of the initial instruction to avoid confusion
        messages = [
            {
                "role": "user",
                "content": self.feedback_prompt(
                    in_context_example=task_in_context_example[
                        task_in_context_example.find("Task:") :
                    ],  # This is to get rid of the initial instruction to avoid confusion
                    trajectory=trajectory,
                    correct_solution=gt_solution,
                    tool_desc=tool_desc,
                ),
            }
        ]

        # log in yellow
        LOGGER.debug(
            "Feedback Agent Prompt:\n" + "\033[93m" + messages[0]["content"] + "\033[0m"
        )
        lm_output, token_usage = self.call_lm(messages)
        for usage_type, count in token_usage.items():
            state.token_counter["feedback_" + usage_type] += count
        action = self.lm_output_to_action(lm_output, form)
        # log in red
        LOGGER.debug("Feedback Agent Action:\n" + "\033[91m" + action.value + "\033[0m")
        return action

    @backoff.on_exception(
        backoff.expo,
        requests.exceptions.RequestException,
    )
    def call_lm(self, messages):
        # Prepend the prompt with the system message
        data = {
            "model": self.config["model_name"],
            "prompt": f"\n\nHuman: {messages[0]['content']}\n\n",
            "max_tokens_to_sample": self.config.get("max_tokens", 512),
            "temperature": self.config.get("temperature", 0),
            "stop_sequences": self.stop_words,
        }
        assert len(messages) == 1, "message length must be 1"

        response = requests.post(url, headers=headers, json=data)

        if response.status_code == 200:
            pass
        else:
            raise requests.exceptions.RequestException(
                "Request failed with status code:", response.status_code
            )

        return json.loads(response.text)["completion"], {}

mint/agents/claude_feedback_agent.py
xingyaoww/mint-bench
import re
import logging

LOGGER = logging.getLogger("MINT")

from .openai_lm_agent import OpenAILMAgent

from mint.datatypes import State, Action
from mint.prompt import FeedbackPromptTemplate
import openai
import logging
import traceback
from mint.datatypes import Action


class OpenAIFeedbackAgent(OpenAILMAgent):
    def __init__(self, config):
        super().__init__(config)
        # The agent should not generate Assistant msg since it should provide feedback
        self.stop_words = ["\nObservation:", "\nTask:", "\nAssistant:"]
        self.feedback_prompt = FeedbackPromptTemplate()

    def lm_output_to_action(self, lm_output, form) -> Action:
        if form == "textual":
            feedback = lm_output
        elif form == "binary":
            # Find the first sentence (as feedback).
            first_sent = re.findall(r"([^.]*\.)", lm_output)[0]
            if "GOOD" in first_sent:
                feedback = "This is GOOD."
            elif "BAD" in first_sent:
                feedback = "This is BAD."
            else:
                raise ValueError(f"Cannot find GOOD or BAD in feedback: {feedback}")
        return Action(feedback, use_tool=False)

    def act(
        self,
        state: State,
        observation: str,
        form: str,
        gt,
        task_in_context_example: str,
        tool_desc: str,
    ) -> Action:
        try:
            gt_solution = (
                (
                    f"Correct solution (please DO NOT disclose the correct solution to the assistant): {str(gt).strip()}\n"
                )
                if gt
                else "Correct solution (please DO NOT disclose the correct solution to the assistant): NOT GIVEN\n"
            )
            trajectory = (
                "---\n".join(state.history[0]["content"].split("---\n")[2:]) + "\n"
            )
            trajectory += "\n".join([x["content"] for x in state.history[1:]])
            trajectory += "\n" + observation
            trajectory = trajectory[
                trajectory.find("Task:") :
            ]  # Get rid of the initial instruction to avoid confusion
            messages = [
                {
                    "role": "user",
                    "content": self.feedback_prompt(
                        in_context_example=task_in_context_example[
                            task_in_context_example.find("Task:") :
                        ],  # This is to get rid of the initial instruction to avoid confusion
                        trajectory=trajectory,
                        correct_solution=gt_solution,
                        tool_desc=tool_desc,
                    ),
                }
            ]

            # log in yellow
            LOGGER.debug(
                "Feedback Agent Prompt:\n"
                + "\033[93m"
                + messages[0]["content"]
                + "\033[0m"
            )
            lm_output, token_usage = self.call_lm(messages)
            for usage_type, count in token_usage.items():
                state.token_counter["feedback_" + usage_type] += count
            action = self.lm_output_to_action(lm_output, form)
            # log in red
            LOGGER.debug(
                "Feedback Agent Action:\n" + "\033[91m" + action.value + "\033[0m"
            )
            return action
        except openai.error.InvalidRequestError:
            tb = traceback.format_exc()
            return Action(f"", False, error=f"InvalidRequestError\n{tb}")
        except Exception as e:
            tb = traceback.format_exc()
            return Action(f"", False, error=f"Unknown error\n{tb}")

mint/agents/openai_feedback_agent.py
xingyaoww/mint-bench
from .base import LMAgent
import openai
import logging
import traceback
from mint.datatypes import Action
import backoff

LOGGER = logging.getLogger("MINT")


class OpenAILMAgent(LMAgent):
    def __init__(self, config):
        super().__init__(config)
        assert "model_name" in config.keys()

    @backoff.on_exception(
        backoff.fibo,
        # https://platform.openai.com/docs/guides/error-codes/python-library-error-types
        (
            openai.error.APIError,
            openai.error.Timeout,
            openai.error.RateLimitError,
            openai.error.ServiceUnavailableError,
            openai.error.APIConnectionError,
        ),
    )
    def call_lm(self, messages):
        # Prepend the prompt with the system message
        response = openai.ChatCompletion.create(
            model=self.config["model_name"],
            messages=messages,
            max_tokens=self.config.get("max_tokens", 512),
            temperature=self.config.get("temperature", 0),
            stop=self.stop_words,
        )
        return response.choices[0].message["content"], response["usage"]

    def act(self, state):
        messages = state.history
        try:
            lm_output, token_usage = self.call_lm(messages)
            for usage_type, count in token_usage.items():
                state.token_counter[usage_type] += count
            action = self.lm_output_to_action(lm_output)
            return action
        except openai.error.InvalidRequestError:  # mostly due to model context window limit
            tb = traceback.format_exc()
            return Action(f"", False, error=f"InvalidRequestError\n{tb}")
        # except Exception as e:
        #     tb = traceback.format_exc()
        #     return Action(f"", False, error=f"Unknown error\n{tb}")

mint/agents/openai_lm_agent.py
xingyaoww/mint-bench
from .openai_lm_agent import OpenAILMAgent
import openai
import openai.error
import logging
import traceback
from mint.datatypes import Action
import backoff

LOGGER = logging.getLogger("MINT")
# REMEMBER to RUN ALL MODELS


class VLLMAgent(OpenAILMAgent):
    """Inference for open-sourced models with a unified interface with OpenAI's API."""

    def __init__(self, config):
        super().__init__(config)
        assert (
            "openai.api_base" in config.keys()
        ), "missing openai.api_base to connect to server"
        self.api_base = config["openai.api_base"]
        self.api_key = "EMPTY"
        LOGGER.info("remember to openup the server using docs/SERVING.mdh")
        self.stop_words = [
            "\nObservation:",
            "\nExpert feedback:",
            "\nTask:",
            "\n---",
            "\nHuman:",
        ]

    def format_prompt(self, messages):
        """Format messages into a prompt for the model."""
        prompt = ""
        for message in messages:
            if message["role"] == "user":
                prompt += f"\n\nHuman: {message['content']}"
            elif message["role"] == "assistant":
                prompt += f"\n\nAssistant: {message['content']}"
        prompt += "\n\nAssistant:"
        return prompt

    @backoff.on_exception(
        backoff.fibo,
        # https://platform.openai.com/docs/guides/error-codes/python-library-error-types
        (
            openai.error.Timeout,
            openai.error.RateLimitError,
            openai.error.ServiceUnavailableError,
            openai.error.APIConnectionError,
        ),
    )
    def call_lm(self, messages):
        if self.config.get("add_system_message", False):
            messages = self.add_system_message(messages)
            assert messages[0]["role"] == "system"
            # system msg will be formatted by vllm and fastchat, so no need to format here
        else:
            messages = [
                {"role": "system", "content": ""}
            ] + messages  # add empty system message

        try:
            if self.config["chat_mode"]:
                response = openai.ChatCompletion.create(
                    model=self.config["model_name"],
                    messages=messages,
                    max_tokens=self.config.get("max_tokens", 512),
                    temperature=self.config.get("temperature", 0),
                    stop=self.stop_words,
                    api_base=self.api_base,
                    api_key=self.api_key,
                )
                resp_str = response.choices[0].message["content"]

            else:
                prompt = self.format_prompt(messages)
                response = openai.Completion.create(
                    model=self.config["model_name"],
                    prompt=prompt,
                    max_tokens=self.config.get("max_tokens", 512),
                    temperature=self.config.get("temperature", 0),
                    stop=self.stop_words,
                    api_base=self.api_base,
                    api_key=self.api_key,
                )
                resp_str = response.choices[0].text

        except openai.error.APIError as e:
            # This is special handling for FastChat Library
            # and is actually unrelated to the OpenAI API
            error_message = e.args[0]
            # Invalid response object from API: '{"object":"error","message":"This model\'s maximum context length is 4096 tokens. However, you requested 4169 tokens (3657 in the messages, 512 in the completion). Please reduce the length of the messages or completion.","type":"invalid_request_error","param":null,"code":null}' (HTTP response code was 400))
            if "maximum context length" in error_message:
                raise openai.error.InvalidRequestError(e.args[0], "")
            else:
                raise e
        resp_str = resp_str.rstrip()  # remove trailing spaces (usually caused by llama)
        return resp_str, response["usage"]

mint/agents/vllm_agent.py
xingyaoww/mint-bench
from .openai_lm_agent import OpenAILMAgent
import openai
import openai.error
import logging
import traceback
from mint.datatypes import State, Action
import backoff
import re
from mint.prompt import FeedbackPromptTemplate

LOGGER = logging.getLogger("MINT")
# REMEMBER to RUN ALL MODELS


class VLLMFeedbackAgent(OpenAILMAgent):
    """Inference for open-sourced models with a unified interface with OpenAI's API."""

    def __init__(self, config):
        super().__init__(config)
        assert (
            "openai.api_base" in config.keys()
        ), "missing openai.api_base to connect to server"
        self.api_base = config["openai.api_base"]
        self.api_key = "EMPTY"
        LOGGER.info("remember to openup the server following docs/SERVING.md")

        if "override" in config:
            self.assistant_keyword = config["override"].get("assistant", "Assistant")
            self.human_keyword = config["override"].get("human", "Human")
            LOGGER.info(
                f"Overriding assistant/human keyword for the [Feedback Model] to {self.assistant_keyword}/{self.human_keyword}"
            )
        else:
            self.assistant_keyword = "Assistant"
            self.human_keyword = "Human"

        self.stop_words = [
            "\nObservation:",
            "\nTask:",
            f"\n{self.assistant_keyword}:",
            f"\n{self.human_keyword}:"
        ]

        self.feedback_prompt = FeedbackPromptTemplate()

    def lm_output_to_action(self, lm_output, form) -> Action:
        if form == "textual":
            feedback = lm_output
        elif form == "binary":
            # Find the first sentence (as feedback).
            first_sent = re.findall(r"([^.]*\.)", lm_output)[0]
            if "GOOD" in first_sent:
                feedback = "This is GOOD."
            elif "BAD" in first_sent:
                feedback = "This is BAD."
            else:
                raise ValueError(f"Cannot find GOOD or BAD in feedback: {feedback}")
        return Action(feedback, use_tool=False)

    def act(
        self,
        state: State,
        observation: str,
        form: str,
        gt,
        task_in_context_example: str,
        tool_desc: str,
    ) -> Action:
        try:
            gt_solution = (
                (
                    f"Correct solution (please DO NOT disclose the correct solution to the assistant): {str(gt).strip()}\n"
                )
                if gt
                else "Correct solution (please DO NOT disclose the correct solution to the assistant): NOT GIVEN\n"
            )
            trajectory = (
                "---\n".join(state.history[0]["content"].split("---\n")[2:]) + "\n"
            )
            trajectory += "\n".join([x["content"] for x in state.history[1:]])
            trajectory += "\n" + observation
            trajectory = trajectory[
                trajectory.find("Task:") :
            ]  # Get rid of the initial instruction to avoid confusion

            feedback_prompt = self.feedback_prompt(
                in_context_example=task_in_context_example[
                    task_in_context_example.find("Task:") :
                ],  # This is to get rid of the initial instruction to avoid confusion
                trajectory=trajectory,
                correct_solution=gt_solution,
                tool_desc=tool_desc,
            )
            # try to map the assistant/human keyword to avoid collision with the LLM's keyword
            feedback_prompt = feedback_prompt.replace("Assistant:", f"{self.assistant_keyword}:")
            feedback_prompt = feedback_prompt.replace("Human:", f"{self.human_keyword}:")
            messages = [
                {
                    "role": "user",
                    "content": feedback_prompt
                }
            ]

            # log in yellow
            LOGGER.debug(
                "Feedback Agent Prompt:\n"
                + "\033[93m"
                + messages[0]["content"]
                + "\033[0m"
            )
            lm_output, token_usage = self.call_lm(messages)
            for usage_type, count in token_usage.items():
                state.token_counter["feedback_" + usage_type] += count
            action = self.lm_output_to_action(lm_output, form)
            # log in red
            LOGGER.debug(
                "Feedback Agent Action:\n" + "\033[91m" + action.value + "\033[0m"
            )
            return action
        except openai.error.InvalidRequestError:
            tb = traceback.format_exc()
            return Action(f"", False, error=f"InvalidRequestError\n{tb}")
        except Exception as e:
            tb = traceback.format_exc()
            return Action(f"", False, error=f"Unknown error\n{tb}")

    def format_prompt(self, messages):
        """Format messages into a prompt for the model."""
        prompt = ""
        for message in messages:
            if message["role"] == "user":
                prompt += f"\n\nHuman: {message['content']}"
            elif message["role"] == "assistant":
                prompt += f"\n\nAssistant: {message['content']}"
        return prompt

    @backoff.on_exception(
        backoff.fibo,
        # https://platform.openai.com/docs/guides/error-codes/python-library-error-types
        (
            openai.error.Timeout,
            openai.error.RateLimitError,
            openai.error.ServiceUnavailableError,
            openai.error.APIConnectionError,
        ),
    )
    def call_lm(self, messages):
        if self.config.get("add_system_message", False):
            messages = self.add_system_message(messages)
            assert messages[0]["role"] == "system"
            # system msg will be formatted by vllm and fastchat, so no need to format here
        else:
            messages = [
                {"role": "system", "content": ""}
            ] + messages  # add empty system message

        try:
            if self.config["chat_mode"]:
                response = openai.ChatCompletion.create(
                    model=self.config["model_name"],
                    messages=messages,
                    max_tokens=self.config.get("max_tokens", 512),
                    temperature=self.config.get("temperature", 0),
                    stop=self.stop_words,
                    api_base=self.api_base,
                    api_key=self.api_key,
                )
                resp_str = response.choices[0].message["content"]

            else:
                prompt = self.format_prompt(messages)
                response = openai.Completion.create(
                    model=self.config["model_name"],
                    prompt=prompt,
                    max_tokens=self.config.get("max_tokens", 512),
                    temperature=self.config.get("temperature", 0),
                    stop=self.stop_words + ["\nExpert feedback:"],
                    api_base=self.api_base,
                    api_key=self.api_key,
                )
                resp_str = response.choices[0].text

        except openai.error.APIError as e:
            # This is special handling for FastChat Library
            # and is actually unrelated to the OpenAI API
            error_message = e.args[0]
            # Invalid response object from API: '{"object":"error","message":"This model\'s maximum context length is 4096 tokens. However, you requested 4169 tokens (3657 in the messages, 512 in the completion). Please reduce the length of the messages or completion.","type":"invalid_request_error","param":null,"code":null}' (HTTP response code was 400))
            if "maximum context length" in error_message:
                raise openai.error.InvalidRequestError(e.args[0], "")
            else:
                raise e
        resp_str = resp_str.rstrip()  # remove trailing spaces (usually caused by llama)
        return resp_str, response["usage"]

mint/agents/vllm_feedback_agent.py
xingyaoww/mint-bench
# all configs are flexible to adapt with this config_variables.py
# if you want to change data output dir or dataset dir, please change DATA_OUTPUTS_DIR and DATASET_DIR
# if you want to add more configs, please:
# (1) if you want to add tools, please add to TASK_TYPE_TO_TOOL_IMPORT
# (2) if you want to add tasks, please add to TASK_INFO_MAP
# (3) if you want to add agent models, please add to agent_model_infos
# (4) if you want to add feedback models and configs, please add to feedback_model_infos and feedback_types
# (5) if you want to add env configs, please add to env_configs

DATA_OUTPUTS_DIR = "data/outputs"
DATASET_DIR = "data/processed"

TASK_TYPE_TO_TOOL_IMPORT = {
    "reasoning": [("mint.tools.wikipedia_search", "WikipediaQueryRun")],
    "decision_making": [],
    "code_generation": [],
}

TASK_INFO_MAP = {
    # === Reasoning ===
    "theoremqa": {"class": "TheoremqaTask", "type": "reasoning"},
    "gsm8k": {"class": "ReasoningTask", "type": "reasoning"},
    "hotpotqa": {"class": "ReasoningTask", "type": "reasoning"},
    "mmlu": {"class": "MultipleChoiceTask", "type": "reasoning"},
    "math": {"class": "ReasoningTask", "type": "reasoning"},
    # === Code Generation ===
    "mbpp": {
        "class": "MBPPTask",
        "type": "code_generation",
    },
    "humaneval": {
        "class": "HumanEvalTask",
        "type": "code_generation",
    },
    # === Decision Making ===
    "alfworld": {"class": "AlfWorldTask", "type": "decision_making"},
}

FEEDBACK_CONFIG = {
    "feedback_agent_config": {
        "chat_mode": True,
        "max_tokens": 1024,
        "temperature": 0.0,
        "stop": ["\nQ:"],
    }
}

EVALUATED_MODEL_LIST = [
    {
        "agent_class": "OpenAILMAgent",
        "config": {
            "model_name": "gpt-3.5-turbo-0613",
            "chat_mode": True,
            "max_tokens": 512,
            "temperature": 0.0,
        },
    },
    # {
    #     "agent_class": "OpenAILMAgent",
    #     "config": {
    #         "model_name": "gpt-3.5-turbo-16k-0613",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #     },
    # },
    # {
    #     "agent_class": "OpenAILMAgent",
    #     "config": {
    #         "model_name": "gpt-4-0613",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #     },
    # },
    # Bard
    # {
    #     "agent_class": "BardLMAgent",
    #     "config": {
    #         "model_name": "chat-bison-001",
    #         "add_system_message": False,
    #         "temperature": 0.0,
    #         "candidate_count": 1,
    #     },
    # },
    # LLaMA2
    # {
    #     "agent_class": "LLaMA2LMAgent",
    #     "config": {
    #         "model_name": "Llama-2-7b-chat-hf",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    # },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "Llama-2-7b-hf",
    #         "chat_mode": False,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "Llama-2-13b-hf",
    #         "chat_mode": False,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "Llama-2-70b-hf",
    #         "chat_mode": False,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "Llama-2-7b-chat-hf",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "Llama-2-13b-chat-hf",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "Llama-2-70b-chat-hf",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "llama-2-lemur-70b-chat-v1",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "llama-2-lemur-70b-v1",
    #         "chat_mode": False,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "vicuna-7b-v1.5",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "vicuna-13b-v1.5",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "vicuna-7b-v1.5-16k",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "vicuna-13b-v1.5-16k",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "CodeLlama-7b-hf",
    #         "chat_mode": False,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "CodeLlama-13b-hf",
    #         "chat_mode": False,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "CodeLlama-34b-hf",
    #         "chat_mode": False,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "CodeLlama-7b-Instruct-hf",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "CodeLlama-13b-Instruct-hf",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "VLLMAgent",
    #     "config": {
    #         "model_name": "CodeLlama-34b-Instruct-hf",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #         "openai.api_base": "<YOUR_API_BASE>",
    #         "add_system_message": False,
    #     },
    # },
    # {
    #     "agent_class": "ClaudeLMAgent",
    #     "config": {
    #         "model_name": "claude-2",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #     },
    # },
    # {
    #     "agent_class": "ClaudeLMAgent",
    #     "config": {
    #         "model_name": "claude-instant-1",
    #         "chat_mode": True,
    #         "max_tokens": 512,
    #         "temperature": 0.0,
    #     },
    # },
]

FEEDBACK_PROVIDER_LIST = [
    {
        "agent_class": "None",
        "model_name": "None",
    },
    {
        "agent_class": "OpenAIFeedbackAgent",
        "model_name": "gpt-4-0613",
    },
    # {
    #     "agent_class": "OpenAIFeedbackAgent",
    #     "model_name": "gpt-3.5-turbo-16k-0613",
    # },
    # {
    #     "agent_class": "ClaudeFeedbackAgent",
    #     "model_name": "claude-2",
    # },
    # {
    #     "agent_class": "ClaudeFeedbackAgent",
    #     "model_name": "claude-instant-1",
    # },
    # {
    #     "agent_class": "VLLMFeedbackAgent",
    #     "model_name": "Llama-2-70b-chat-hf",
    #     "openai.api_base": "<YOUR_API_BASE>",
    #     "chat_mode": True,
    # },
    # {
    #     "agent_class": "VLLMFeedbackAgent",
    #     "model_name": "CodeLlama-34b-hf",
    #     "openai.api_base": "<YOUR_API_BASE>",
    #     "chat_mode": False,
    # },
    # {
    #     "agent_class": "VLLMFeedbackAgent",
    #     "model_name": "Llama-2-70b-hf",
    #     "openai.api_base": "<YOUR_API_BASE>",
    #     "chat_mode": False,
    # },
    # {
    #     "agent_class": "VLLMFeedbackAgent",
    #     "model_name": "CodeLlama-34b-Instruct-hf",
    #     "openai.api_base": "<YOUR_API_BASE>",
    #     "chat_mode": False,
    # },
]

FEEDBACK_TYPES = [
    {"pseudo_human_feedback": "no_GT", "feedback_form": "textual"},
    # {"pseudo_human_feedback": "no_GT", "feedback_form": "binary"},
    # {"pseudo_human_feedback": "GT", "feedback_form": "binary"},
    # {"pseudo_human_feedback": "GT", "feedback_form": "textual"},
]

ENV_CONFIGS = [
    {
        "max_steps": 5,
        "use_tools": True,
        "max_propose_solution": 2,
        "count_down": True,
    },
    {
        "max_steps": 4,
        "use_tools": True,
        "max_propose_solution": 2,
        "count_down": True,
    },
    {
        "max_steps": 3,
        "use_tools": True,
        "max_propose_solution": 2,
        "count_down": True,
    },
    {
        "max_steps": 2,
        "use_tools": True,
        "max_propose_solution": 2,
        "count_down": True,
    },
    {
        "max_steps": 1,
        "use_tools": True,
        "max_propose_solution": 1,
        "count_down": True,
    },
    # {
    #     "max_steps": 6,
    #     "use_tools": True,
    #     "max_propose_solution": 2,
    #     "count_down": True,
    # },
    # {
    #     "max_steps": 8,
    #     "use_tools": True,
    #     "max_propose_solution": 2,
    #     "count_down": True,
    # },
    # {
    #     "max_steps": 10,
    #     "use_tools": True,
    #     "max_propose_solution": 2,
    #     "count_down": True,
    # }
]

mint/configs/config_variables.py
xingyaoww/mint-bench
import os
import json
import pathlib
from copy import deepcopy
import glob
import shutil

# load variables from config_variables.py
from mint.configs.config_variables import (
    EVALUATED_MODEL_LIST,
    FEEDBACK_PROVIDER_LIST,
    FEEDBACK_TYPES,
    ENV_CONFIGS,
    TASK_INFO_MAP,
    TASK_TYPE_TO_TOOL_IMPORT,
    FEEDBACK_CONFIG,
    DATA_OUTPUTS_DIR,
    DATASET_DIR,
)


def build_path(
    task_name: str,
    agent_model_name: str,
    feedback_type: dict,
    feedback_model_name: str,
    env_config: str,
    prefix: str,
):
    dir_path = prefix
    # 1. first-level dir = agent model name
    dir_path = os.path.join(dir_path, agent_model_name)

    # 2. second-level dir = feedback model name
    dir_path = os.path.join(dir_path, "F=" + feedback_model_name)
    if feedback_model_name != "None":
        # 3. (optional) third-level dir = feedback type
        pseudo_human_feedback = feedback_type["pseudo_human_feedback"]
        feedback_form = feedback_type["feedback_form"]
        dir_path = os.path.join(
            dir_path, f"PHF={pseudo_human_feedback}-{feedback_form}"
        )

    # 4. fourth-level dir = env config
    max_propose_solution = env_config["max_propose_solution"]
    max_steps = env_config["max_steps"]
    use_tools = env_config["use_tools"]
    count_down = env_config["count_down"]
    cur_dir_name = f"max{max_steps}_p{max_propose_solution}"
    if use_tools:
        cur_dir_name += "+tool"
    if count_down:
        cur_dir_name += "+cd"
    dir_path = os.path.join(dir_path, cur_dir_name)

    # 5. fifth-level dir = task type
    dir_path = os.path.join(dir_path, TASK_INFO_MAP[task_name]["type"])

    # 6. sixth-level dir = task name
    dir_path = os.path.join(dir_path, task_name)
    return dir_path


def generate_config_json(
    task_name: str,
    agent_model_info: dict,
    feedback_type: dict,
    feedback_model_info: dict,
    env_config: str,
):
    model_name = agent_model_info["config"]["model_name"]
    if agent_model_info.get("exp_setting", "") == "human-eval-human":
        model_name += "-human-eval-human"
    elif agent_model_info.get("exp_setting", "") == "human-eval-gpt":
        model_name += "-human-eval-gpt"
    elif agent_model_info.get("exp_setting", "") == "human-eval-none":
        model_name += "-human-eval-none"

    output_filepath = build_path(
        task_name=task_name,
        agent_model_name=model_name,
        feedback_type=feedback_type,
        feedback_model_name=feedback_model_info["model_name"],
        env_config=env_config,
        prefix=DATA_OUTPUTS_DIR,
    )

    json_dict = {
        "agent": agent_model_info,
        "task": {
            "task_class": TASK_INFO_MAP[task_name]["class"],
            "task_type": TASK_INFO_MAP[task_name]["type"],
            "tool_imports": TASK_TYPE_TO_TOOL_IMPORT[TASK_INFO_MAP[task_name]["type"]],
        },
        "output_dir": output_filepath,
        "env_config": env_config,
    }
    pathlib.Path(output_filepath).mkdir(parents=True, exist_ok=True)
    with open(output_filepath + "/output.txt", "a") as f:
        pass

    if agent_model_info.get("exp_setting", "") == "human-eval-human":
        json_dict["task"][
            "filepath"
        ] = f"data/human_preference/test_human/{task_name}/" + (
            "test_prompts.json"
            if TASK_INFO_MAP[task_name]["type"]
            in ["reasoning", "code_generation", "vision_language"]
            else ""
        )
    elif agent_model_info.get("exp_setting", "") == "human-eval-gpt":
        json_dict["task"][
            "filepath"
        ] = f"data/human_preference/test_gpt/{task_name}/" + (
            "test_prompts.json"
            if TASK_INFO_MAP[task_name]["type"]
            in ["reasoning", "code_generation", "vision_language"]
            else ""
        )
    elif agent_model_info.get("exp_setting", "") == "human-eval-none":
        json_dict["task"][
            "filepath"
        ] = f"data/human_preference/test_none/{task_name}/" + (
            "test_prompts.json"
            if TASK_INFO_MAP[task_name]["type"]
            in ["reasoning", "code_generation", "vision_language"]
            else ""
        )
    else:
        json_dict["task"]["filepath"] = f"{DATASET_DIR}/{task_name}/" + (
            "test_prompts.json"
            if TASK_INFO_MAP[task_name]["type"]
            in ["reasoning", "code_generation", "vision_language"]
            else ""
        )
    json_dict["feedback_config"] = FEEDBACK_CONFIG
    json_dict["feedback_config"].update(feedback_type)
    json_dict["feedback_config"]["feedback_agent_config"].update(feedback_model_info)

    return json_dict


def build_json_for_all_tasks(
    agent_model_info: dict,
    feedback_type: dict,
    feedback_model_info: dict,
    env_config: dict,
):
    for task_name in TASK_INFO_MAP.keys():
        model_name = agent_model_info["config"]["model_name"]
        if agent_model_info.get("exp_setting", "") == "human-eval-human":
            model_name += "-human-eval-human"
        elif agent_model_info.get("exp_setting", "") == "human-eval-gpt":
            model_name += "-human-eval-gpt"
        elif agent_model_info.get("exp_setting", "") == "human-eval-none":
            model_name += "-human-eval-none"
        output_filepath = (
            build_path(
                task_name=task_name,
                agent_model_name=model_name,
                feedback_type=feedback_type,
                feedback_model_name=feedback_model_info["model_name"],
                env_config=env_config,
                prefix="configs/",
            )
            + ".json"
        )
        json_dict = generate_config_json(
            task_name=task_name,
            agent_model_info=agent_model_info,
            feedback_type=feedback_type,
            feedback_model_info=feedback_model_info,
            env_config=env_config,
        )

        pathlib.Path(output_filepath).parent.mkdir(parents=True, exist_ok=True)
        with open(output_filepath, "w") as f:
            f.write(json.dumps(json_dict, indent=4) + "\n")


if __name__ == "__main__":
    # delete all existing configs
    for file in glob.glob("configs/*"):
        shutil.rmtree(file)

    # With feedback
    for agent_model_info in EVALUATED_MODEL_LIST:
        for feedback_model_info in FEEDBACK_PROVIDER_LIST:
            for env_config in ENV_CONFIGS:

                # If no feedback is provided
                if feedback_model_info["agent_class"] == "None":
                    build_json_for_all_tasks(
                        agent_model_info=agent_model_info,
                        feedback_type={
                            "pseudo_human_feedback": "None",
                            "feedback_form": "None",
                        },
                        feedback_model_info=feedback_model_info,
                        env_config=env_config,
                    )

                # If feedback is provided - generate config for all feedback types
                else:
                    # skip if not k=5
                    if env_config["max_steps"] != 5:
                        continue
                    for feedback_type in FEEDBACK_TYPES:
                        build_json_for_all_tasks(
                            agent_model_info=agent_model_info,
                            feedback_type=feedback_type,
                            feedback_model_info=feedback_model_info,
                            env_config=env_config,
                        )

mint/configs/generate_config.py
xingyaoww/mint-bench
import enum
from dataclasses import dataclass
from typing import List, Dict, Any, Optional, Tuple
from collections import defaultdict


@dataclass(frozen=True)
class Action:
    value: str  # LM returned string for now
    use_tool: bool  # if use_tool == False -> propose answer
    error: Optional[str] = None


class State:
    """This should contains everything needed to continue the conversation.

    For example, the history of the conversation, the current task (success/failure) at each step, etc.
    """

    def __init__(
        self,
        history: List[Dict[str, Any]] = None,
        finished: bool = False,
        success: bool = False,
        latest_output: Dict[str, Any] = None,
        agent_action_count: Dict[str, int] = None,
        terminate_reason: str = None,
    ):
        self.history: List[Dict[str, Any]] = history
        self.finished: bool = finished
        self.success: bool = success
        self.latest_output: Dict[str, Any] = latest_output
        self.agent_action_count: Dict[str, int] = agent_action_count
        self.token_counter: Dict[str, int] = defaultdict(int)
        self.terminate_reason: str = terminate_reason
        self.error: Optional[str] = None

    @property
    def empty(self):
        return len(self.history) == 0

    def to_dict(self) -> Dict[str, Any]:
        return {
            "history": self.history,
            "finished": self.finished,
            "success": self.success,
            "latest_output": self.latest_output,
            "agent_action_count": self.agent_action_count,
            "token_counter": dict(self.token_counter),
            "terminate_reason": self.terminate_reason,
            "error": self.error,
        }


class FeedbackType(enum.Enum):
    FEEDBACK_WITH_GT = "feedback_with_gt"
    FEEDBACK_WO_GT = "feedback_wo_gt"
    NO_FEEDBACK = "no_feedback"


class StepOutput:
    def __init__(
        self,
        observation: str = None,
        feedback: str = "",
        feedback_type: FeedbackType = FeedbackType.NO_FEEDBACK,
        success: bool = False,
        extra: Dict[str, Any] = None,
        turn_info: Tuple[int, int] = None,
    ):
        self.observation: str = observation
        self.feedback: str = feedback
        self.feedback_type: FeedbackType = feedback_type
        self.success: bool = success
        self.extra: Dict[str, Any] = extra
        self.turn_info = turn_info

    def __repr__(self) -> str:
        return self.observation

    def to_str(self) -> str:
        output = "Observation:\n"
        if self.observation is not None:
            output += self.observation + "\n"
        else:
            if self.success == False:
                output += "Your answer is wrong.\n"

        if self.turn_info != None:
            n_steps_left, n_propose_solution_left = self.turn_info
            output += "You have {} steps left and {} chances to propose solution left.\n".format(
                n_steps_left, n_propose_solution_left
            )
            if n_steps_left <= 1:
                output += "You should take the last step to propose a solution.\n"

        if self.feedback_type != FeedbackType.NO_FEEDBACK and self.feedback != "":
            # output += (
            #     f"\nYour {action_type} is commented by an expert "
            #     f"{'without' if self.feedback_type == FeedbackType.FEEDBACK_WO_GT else 'with'}"
            #     " access to the correct answer.\nExpert feedback:\n"
            # )
            # Maybe we don't want to tell the model whether the expert has acces to GT
            output += "\nExpert feedback:\n" + self.feedback + "\n"

        return output

    def to_dict(self) -> Dict[str, Any]:
        return {
            "observation": self.observation,
            "feedback": self.feedback,
            "feedback_type": self.feedback_type.value,
            "success": self.success,
        }

mint/datatypes.py
xingyaoww/mint-bench
from .base import BaseEnv
from .general_env import GeneralEnv
from .alfworld_env import AlfworldEnv

mint/envs/__init__.py
xingyaoww/mint-bench
from mint.envs.general_env import GeneralEnv
from mint.datatypes import StepOutput, Action
from mint.tools import Tool
from mint.tools.alfworld import ALFWORLD_TOOL_CLS
from mint.tasks.alfworld import AlfWorldTask
from typing import Any, Dict, List, Mapping


class AlfworldEnv(GeneralEnv):
    def __init__(
        self,
        task: AlfWorldTask,
        tool_set: List[Tool],
        feedback_config: Dict[str, Any],
        environment_config: Dict[str, Any],
    ):
        self.env = task.env
        self.action_results: List[StepOutput] = []
        self.tool_set = [
            tool_cls(self.env, callback_fn=self.action_results.append)
            for tool_cls in ALFWORLD_TOOL_CLS
        ]  # will be merged with the tool_set in GeneralEnv
        super().__init__(task, tool_set, feedback_config, environment_config)

    def check_task_success(self, *args, **kwargs) -> bool:
        """Check if the task is successful.

        In AlfWorld, we can check the last tool call result to see if the task is successful.
        """
        if len(self.action_results) == 0:
            # If no interaction with AlfWorld, then the task cannot be successful.
            return False
        for result in reversed(self.action_results):
            if result.success:
                return True
        return False

    def handle_tool_call(self, action: Action):
        obs = super().handle_tool_call(action)
        # since in alfworld, making the correct tool call is enough to finish the task
        # so we need to check success and update the state here
        self.state.success = self.check_task_success()
        if self.state.success:
            self.state.finished = True
            self.state.terminate_reason = "task_success"
        return obs

mint/envs/alfworld_env.py
xingyaoww/mint-bench
from abc import ABC, abstractmethod
from mint.datatypes import State, Action


class BaseEnv(ABC):
    @abstractmethod
    def step(self, action: Action) -> State:
        pass

    @abstractmethod
    def reset(self) -> State:
        pass

mint/envs/base.py
xingyaoww/mint-bench
import re
import logging
import traceback
from typing import Any, Dict, List, Mapping, Tuple, Optional

LOGGER = logging.getLogger("MINT")

from mint import agents
from mint.envs.base import BaseEnv
from mint.prompt import ToolPromptTemplate
from mint.datatypes import State, Action, StepOutput, FeedbackType
from mint.tools import Tool, get_toolset_description
from mint.tasks import Task
from mint.tools.python_tool import PythonREPL
from mint.utils.exception import ParseError


INVALID_INPUT_MESSAGE = (
    "I don't understand your input. \n"
    "If you want to execute code, please use <execute> YOUR_CODE_HERE </execute>.\n"
    "If you want to give me an answer, please use <solution> YOUR_SOLUTION_HERE </solution>.\n"
    "For example: The answer to the question is <solution> 42 </solution>. \n"
)


class GeneralEnv(BaseEnv):
    def __init__(
        self,
        task: Task,
        tool_set: List[Tool],
        feedback_config: Dict[str, Any],
        environment_config: Dict[str, Any],
    ):
        self.task: Task = task
        self.tool_set: List[Tool] = tool_set + getattr(self, "tool_set", [])

        self.state = State()
        self.config = environment_config

        # Feedback
        self.feedback_config = feedback_config
        feedback_agent_config = feedback_config["feedback_agent_config"]
        if feedback_config["pseudo_human_feedback"] in ["GT", "no_GT"]:
            self.feedback_agent: agents = getattr(
                agents, feedback_agent_config["agent_class"]
            )(feedback_agent_config)
        else:
            self.feedback_agent = None
        if self.feedback_config["pseudo_human_feedback"] == "None":
            self.feedback_type = FeedbackType.NO_FEEDBACK
        elif self.feedback_config["pseudo_human_feedback"] == "no_GT":
            self.feedback_type = FeedbackType.FEEDBACK_WO_GT
        elif self.feedback_config["pseudo_human_feedback"] == "GT":
            self.feedback_type = FeedbackType.FEEDBACK_WITH_GT
        else:
            raise ValueError(
                f"Invalid feedback type {self.feedback_config['pseudo_human_feedback']}"
            )

        self.env_outputs: List[StepOutput] = []
        LOGGER.info(
            f"{len(self.tool_set)} tools loaded: {', '.join([t.name for t in self.tool_set])}"
        )
        self.python_repl = PythonREPL(
            name_to_func_mapping={tool.name: tool.__call__ for tool in self.tool_set}
        )

    def parse_action(self, action: Action) -> Tuple[str, Dict[str, Any]]:
        """Define the parsing logic."""
        lm_output = "\n" + action.value + "\n"
        output = {}
        try:
            if not action.use_tool:
                answer = "\n".join(
                    [
                        i.strip()
                        for i in re.findall(
                            r"<solution>(.*?)</solution>", lm_output, re.DOTALL
                        )
                    ]
                )
                if answer == "":
                    raise ParseError("No answer found.")
                output["answer"] = answer
            else:
                env_input = "\n".join(
                    [
                        i.strip()
                        for i in re.findall(
                            r"<execute>(.*?)</execute>", lm_output, re.DOTALL
                        )
                    ]
                )
                if env_input == "":
                    raise ParseError("No code found.")
                output["env_input"] = env_input
        except Exception as e:
            raise ParseError(e)
        return output

    def get_feedback(self, observation: str) -> Tuple[str, FeedbackType]:
        if self.feedback_type == FeedbackType.NO_FEEDBACK:
            return ""
        elif self.feedback_type == FeedbackType.FEEDBACK_WO_GT:
            gt = None
        else:
            gt = self.task.reference

        feedback = self.feedback_agent.act(
            self.state,
            observation=observation,
            form=self.feedback_config["feedback_form"],
            gt=gt,
            task_in_context_example=self.task.in_context_example(
                use_tool=self.config["use_tools"],
                with_feedback=True,
            ),
            tool_desc=get_toolset_description(self.tool_set),
        )
        return feedback.value

    def check_task_success(self, answer: str) -> bool:
        LOGGER.info(f"STUDENT ANSWER: [{answer}]")
        LOGGER.info(f"REFERENCE ANSWER: [{self.task.reference}]")
        return self.task.success(answer)

    def log_output(self, output: StepOutput) -> None:
        if self.state.finished:
            return
        content = output.to_str()
        self.state.history.append({"role": "user", "content": content})
        self.state.latest_output = output.to_dict()
        self.state.latest_output["content"] = content

    def handle_tool_call(self, action: Action) -> str:
        """Use tool to obtain "observation."""
        try:
            parsed = self.parse_action(action)
            env_input = parsed["env_input"]
            obs = self.python_repl(env_input).strip()
            self.env_outputs.append(StepOutput(observation=obs))
            self.state.agent_action_count["use_tool"] += 1
            return obs
        except ParseError:
            self.state.agent_action_count["invalid_action"] += 1
            return INVALID_INPUT_MESSAGE
        except Exception as e:
            error_traceback = traceback.format_exc()
            return f"{error_traceback}"

    def handle_propose_solution(self, action: Action) -> Optional[str]:
        """Propose answer to check the task success.

        It might set self.state.finished = True if the task is successful.
        """
        self.state.agent_action_count["propose_solution"] += 1
        try:
            parsed = self.parse_action(action)
            task_success = self.check_task_success(parsed["answer"])
            if task_success:
                self.state.finished = True
                self.state.success = True
                self.state.terminate_reason = "task_success"
                # NOTE: should not return the function now, because we need to log the output
                # Set state.finished = True will terminate the episode
        except ParseError:
            return INVALID_INPUT_MESSAGE
        except Exception as e:
            error_traceback = traceback.format_exc()
            return f"{error_traceback}"

    def check_max_iteration(self):
        """Check if the agent has reached the max iteration limit.

        It might set self.state.finished = True if the agent has reached the max iteration limit.
        """
        if self.state.finished:
            # ignore if the episode is already finished (e.g., task success)
            return

        if (
            # propose solution > max output solution
            self.state.agent_action_count["propose_solution"]
            >= self.config["max_propose_solution"]
        ):
            self.state.finished = True
            self.state.success = False
            self.state.terminate_reason = "max_propose_steps"
        elif (
            # (propose_solution + use_tool) > max iteration limit
            sum(self.state.agent_action_count.values())
            >= self.config["max_steps"]
        ):
            self.state.finished = True
            self.state.success = False
            self.state.terminate_reason = "max_steps"

    def step(self, action: Action, loaded=None) -> State:
        assert (
            not self.state.finished
        ), "Expecting state.finished == False for env.step()."

        # Update state by logging the action
        if action.value:
            assistant_action = (
                "Assistant:\n" + action.value
                if not action.value.lstrip().startswith("Assistant:")
                else action.value
            )
            self.state.history.append(
                {"role": "assistant", "content": assistant_action + "\n"}
            )

        if action.error:
            # Check if error (usually hit the max length)
            observation = f"An error occurred. {action.error}"
            self.state.finished = True
            self.state.success = False
            self.state.error = action.error
            self.state.terminate_reason = "error"
            LOGGER.error(f"Error:\n{action.error}")
        elif action.use_tool:
            observation = self.handle_tool_call(action)
        else:
            # It might set self.state.finished = True if the task is successful.
            observation = self.handle_propose_solution(action)

        # Check if the agent has reached the max iteration limit.
        # If so, set self.state.finished = True
        # This corresponds to a no-op if the episode is already finished
        self.check_max_iteration()

        # record the turn info
        if self.config["count_down"]:
            turn_info = (
                self.config["max_steps"] - sum(self.state.agent_action_count.values()),
                self.config["max_propose_solution"]
                - self.state.agent_action_count["propose_solution"],
            )
        else:
            turn_info = None

        # Get feedback if the episode is not finished
        if loaded != None:
            feedback = loaded["feedback"]
            LOGGER.info(f"Loaded feedback: {feedback}")
        elif not self.state.finished:
            # This is the output without feedback
            # use to generate an observation for feedback agent
            tmp_output = StepOutput(
                observation=observation,
                success=self.state.success,
                turn_info=turn_info,
            )
            feedback = self.get_feedback(observation=tmp_output.to_str())
        else:
            feedback = ""

        # Log the output to state regardless of whether the episode is finished
        output = StepOutput(
            observation=observation,
            feedback=feedback,
            feedback_type=self.feedback_type,
            success=self.state.success,
            turn_info=turn_info,
        )

        self.log_output(output)
        return self.state

    def reset(self) -> State:
        use_tool: bool = self.config["use_tools"]
        if use_tool and len(self.tool_set) > 0:
            LOGGER.warning(
                (
                    "No tool is provided when use_tools is True.\n"
                    "Ignore this if you are running code generation."
                )
            )

        user_prompt = ToolPromptTemplate(use_tool=use_tool)(
            max_total_steps=self.config["max_steps"],
            max_propose_solution=self.config["max_propose_solution"],
            tool_desc=get_toolset_description(self.tool_set),
            in_context_example=self.task.in_context_example(
                use_tool=use_tool,
                with_feedback=self.feedback_type != FeedbackType.NO_FEEDBACK,
            ),
            task_prompt="Task:\n" + self.task.prompt,
        )
        self.state.history = [{"role": "user", "content": user_prompt}]
        self.state.latest_output = {"content": user_prompt}
        self.state.agent_action_count = {
            "propose_solution": 0,
            "use_tool": 0,
            "invalid_action": 0,
        }

        if use_tool:
            # reset tool set
            for tool in self.tool_set:
                tool.reset()
        return self.state

mint/envs/general_env.py
xingyaoww/mint-bench
from mint.envs import GeneralEnv, AlfworldEnv
from mint.datatypes import Action, State
from mint.tasks import AlfWorldTask
from mint.tools import Tool
import mint.tasks as tasks
import mint.agents as agents
import logging
import os
import json
import pathlib
import importlib
import argparse
from typing import List, Dict, Any
from tqdm import tqdm
from tqdm.contrib.logging import logging_redirect_tqdm

# Configure logging settings
logging.basicConfig(
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
LOGGER = logging.getLogger("MINT")


def interactive_loop(
    task: tasks.Task,
    agent: agents.LMAgent,
    tools: List[Tool],
    feedback_config: Dict[str, Any],
    env_config: Dict[str, Any],
    interactive_mode: bool = False,
):
    if isinstance(task, AlfWorldTask):
        LOGGER.info("loading Alfworld Env")
        env = AlfworldEnv(task, tools, feedback_config, env_config)
    else:
        env = GeneralEnv(task, tools, feedback_config, env_config)
    state: State = env.reset()

    init_msg = state.latest_output['content']
    if interactive_mode:
        # omit in-context example
        splited_msg = init_msg.split("---")
        init_msg = splited_msg[0] + "== In-context Example Omitted ==" + splited_msg[2]

    LOGGER.info(f"\nUser: \n\033[94m{state.latest_output['content']}\033[0m")

    num_steps = 0

    if task.loaded_history is not None:
        for turn in task.loaded_history:
            action = agent.lm_output_to_action(turn["lm_output"])
            LOGGER.info(
                f"\nLoaded LM Agent Action:\n\033[92m{action.value}\033[0m")
            state = env.step(action, loaded=turn)
            LOGGER.info(
                "\033[1m" + "User:\n" + "\033[0m" +
                f"\033[94m{state.latest_output['content']}\033[0m"
            )
            num_steps += 1

    while not state.finished:
        # agent act
        if interactive_mode:
            to_continue = "n"
            while to_continue not in ["y", "Y"]:
                to_continue = input("\n> Continue? (y/n) ")

        action: Action = agent.act(state)
        # color the action in green
        # LOGGER.info(f"\nLM Agent Action:\n\033[92m{action.value}\033[0m")
        LOGGER.info(
            f"\n\033[1m" + "LM Agent Action:\n" + "\033[0m" +
            f"\n\033[92m{action.value}\033[0m"
        )
        # environment step
        state: State = env.step(action)
        # color the state in blue
        if not state.finished:
            user_msg = state.latest_output['content']
            if "Expert feedback:" in user_msg:
                try:
                    splited = user_msg.split("Expert feedback:")
                    obs, feedback = splited[0], "".join(splited[1:])
                except Exception as e:
                    LOGGER.info(f"Error: {e}")
                    LOGGER.info(f"User message: {user_msg}")
                    LOGGER.exception(e)
                    raise e

                feedback = "Expert feedback:" + feedback
                # color the observation in blue & feedback in red
                LOGGER.info(
                    "\n" +
                    "\033[1m" + "User:\n" + "\033[0m" +
                    f"\033[94m{obs}\033[0m" + "\n" 
                    + f"\033[93m{feedback}\033[0m" + "\n"
                )
            else:
                # color the observation in blue
                LOGGER.info(
                    "\n" +
                    "\033[1m" + "User:\n" + "\033[0m" +
                    f"\033[94m{user_msg}\033[0m" + "\n"
                )
        num_steps += 1

    LOGGER.info(
        f"Task finished in {num_steps} steps. Success: {state.success}"
    )

    return state


def main(args: argparse.Namespace):
    with open(args.exp_config) as f:
        exp_config: Dict[str, Any] = json.load(f)

    DEFAULT_FEEDBACK_CONFIG = exp_config["feedback_config"]
    DEFAULT_ENV_CONFIG = exp_config["env_config"]

    LOGGER.info(f"Experiment config: {exp_config}")

    # initialize all the tasks
    task_config: Dict[str, Any] = exp_config["task"]
    task_class: tasks.Task = getattr(tasks, task_config["task_class"])
    todo_tasks, n_tasks = task_class.load_tasks(task_config["filepath"])

    # initialize the agent
    agent_config: Dict[str, Any] = exp_config["agent"]
    agent: agents.LMAgent = getattr(agents, agent_config["agent_class"])(
        agent_config["config"]
    )

    # initialize the feedback agent (if exist)
    feedback_config: Dict[str, Any] = exp_config.get(
        "feedback", DEFAULT_FEEDBACK_CONFIG
    )

    # initialize all the tools
    tools: List[Tool] = [
        getattr(importlib.import_module(module), class_name)()
        for module, class_name in task_config["tool_imports"]
    ]

    env_config: Dict[str, Any] = exp_config.get(
        "environment", DEFAULT_ENV_CONFIG)

    pathlib.Path(exp_config["output_dir"]).mkdir(parents=True, exist_ok=True)
    if args.interactive:
        output_path = os.path.join(
            exp_config["output_dir"], "results.interactive.jsonl")
    else:
        output_path = os.path.join(exp_config["output_dir"], "results.jsonl")

    done_task_id = set()
    if os.path.exists(output_path):
        with open(output_path) as f:
            for line in f:
                task_id = json.loads(line)["task"].get("task_id", "")
                if task_id == "":
                    task_id = json.loads(line)["task"].get("id", "")
                done_task_id.add(task_id)
        LOGGER.info(
            f"Existing output file found. {len(done_task_id)} tasks done.")

    if len(done_task_id) == n_tasks:
        LOGGER.info("All tasks done. Exiting.")
        return

    # run the loop for all tasks
    logging.info(f"Running interactive loop for {n_tasks} tasks.")
    n_tasks = n_tasks - len(done_task_id)  # only run the remaining tasks
    with open(output_path, "a") as f, logging_redirect_tqdm():
        pbar = tqdm(total=n_tasks)
        for i, task in enumerate(todo_tasks):
            # Only test 10 tasks in debug mode
            if args.debug and i == 3:
                break

            # skip done tasks
            if task.task_id in done_task_id:
                continue

            state = interactive_loop(
                task, agent, tools, feedback_config, env_config, args.interactive
            )
            if not os.path.exists(exp_config["output_dir"]):
                os.makedirs(exp_config["output_dir"])
            f.write(
                json.dumps({"state": state.to_dict(),
                           "task": task.to_dict()}) + "\n"
            )
            f.flush()  # make sure the output is written to file
            pbar.update(1)
        pbar.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Run the interactive loop.")
    parser.add_argument(
        "--exp_config",
        type=str,
        default="./configs/gpt-3.5-turbo-0613/F=gpt-3.5-turbo-16k-0613/PHF=GT-textual/max5_p2+tool+cd/reasoning/scienceqa.json",
        help="Config of experiment.",
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Whether to run in debug mode (10 ex per task).",
    )
    parser.add_argument(
        "--interactive",
        action="store_true",
        help="Whether to run in interactive mode for demo purpose.",
    )
    args = parser.parse_args()
    LOGGER.setLevel(logging.DEBUG if args.debug else logging.INFO)
    main(args)

mint/main.py
xingyaoww/mint-bench
import os
from mint.utils import load_file

PROMPT_DIR = os.path.join(os.path.dirname(__file__), "templates")
TEMPLATE_WITH_TOOL = load_file(os.path.join(PROMPT_DIR, "template_with_tool.txt"))
TEMPLATE_WITHOUT_TOOL = load_file(os.path.join(PROMPT_DIR, "template_without_tool.txt"))
TEMPLATE_FEEDBACK_AGENT = load_file(
    os.path.join(PROMPT_DIR, "template_feedback_agent.txt")
)


class PromptTemplate:
    """A prompt template."""

    def __init__(self, template: str):
        self.template: str = template

    def __call__(self, **kwargs) -> str:
        return self.template.format(**kwargs)


class ToolPromptTemplate(PromptTemplate):
    def __init__(self, use_tool: bool):
        if use_tool:
            template = TEMPLATE_WITH_TOOL
        else:
            template = TEMPLATE_WITHOUT_TOOL
        super().__init__(template)


class FeedbackPromptTemplate(PromptTemplate):
    def __init__(self):
        super().__init__(TEMPLATE_FEEDBACK_AGENT)

mint/prompt/__init__.py
xingyaoww/mint-bench
from .base import Task
from .reasoning import ReasoningTask, MultipleChoiceTask, TheoremqaTask
from .codegen import CodeGenTask, HumanEvalTask, MBPPTask
from .alfworld import AlfWorldTask

mint/tasks/__init__.py
xingyaoww/mint-bench
import os
import yaml
import logging
from typing import Iterable
import alfworld
import alfworld.agents.environment as envs
import json
import os

LOGGER = logging.getLogger("MINT")

from mint.tasks.base import Task

PREFIXES = {
    "pick_and_place": "put",
    "pick_clean_then_place": "clean",
    "pick_heat_then_place": "heat",
    "pick_cool_then_place": "cool",
    "look_at_obj": "examine",
    "pick_two_obj": "puttwo",
}


class AlfWorldTask(Task):
    """Alfworld task instance."""

    task_name = "alfworld"

    def __init__(
        self,
        id: str,
        prompt: str,
        reference: str,
        env: envs.AlfredTWEnv,
        task_type: str,
        **kwargs,
    ):
        self.task_name = f"alfworld/{task_type}"  # used to load the correct ICL example
        super().__init__(**kwargs)
        self._id = id
        self._prompt = prompt.strip()
        self._reference = reference
        self.metadata["task_type"] = task_type

        # NOTE: AlfWorld is different -> no reference solution
        self._env = env

    @property
    def env(self) -> envs.AlfredTWEnv:
        """Stateful environment of the task.

        Specific for AlfWorld.
        """
        return self._env

    def success(self, solution: str) -> bool:
        """This checks whether the given solution can complete the current task."""
        # Task success is checked at the environment level, not at the solution level.
        raise NotImplementedError

    @classmethod
    def load_tasks(cls, path: str = "./data/processed/alfworld") -> Iterable["Task"]:
        """Load alfworld data and prompts from a directory."""

        with open(os.path.join(path, "base_config.yaml")) as f:
            config = yaml.safe_load(f)

        # Split following ReAct
        # https://github.com/ysymyth/ReAct/blob/master/alfworld.ipynb
        split = "eval_out_of_distribution"
        env = getattr(alfworld.agents.environment, config["env"]["type"])(
            config, train_eval=split
        )
        assert isinstance(env, alfworld.agents.environment.AlfredTWEnv)
        env = env.init_env(batch_size=1)

        N_TASKS = 134  # as hardcoded in ReAct

        def generator():
            loaded_history_path = os.path.join(
                os.path.dirname(path), "test_prompts.json"
            )
            if os.path.exists(loaded_history_path):
                loaded_histories = [
                    json.loads(i) for i in open(loaded_history_path).readlines()
                ]
                loaded_histories = {
                    i["prompt"].strip(): i["loaded_history"] for i in loaded_histories
                }
            else:
                loaded_histories = None

            for _ in range(N_TASKS):
                ob, info = env.reset()
                ob = "\n".join(ob[0].split("\n\n")[1:])
                game_file = info["extra.gamefile"][0]
                gt_reference_file = os.path.join(
                    os.path.dirname(game_file), "gt_traj.txt"
                )

                with open(gt_reference_file) as f:
                    gt_reference = f.read()

                name = "/".join(game_file.split("/")[-3:-1])
                LOGGER.info(f"Loaded task {name} from {split} split.")

                task_type = None
                for _, (k, v) in enumerate(PREFIXES.items()):
                    if name.startswith(k):
                        task_type = v
                        break
                assert task_type is not None, f"Task type not found for {name}"

                prompt = "Interact with a household to solve a task. \n" + ob
                if loaded_histories == None:
                    yield cls(
                        id=name,
                        prompt=prompt,
                        reference=gt_reference,
                        env=env,
                        task_type=task_type,
                        loaded_history=None,
                    )
                else:
                    if prompt.strip() in loaded_histories:
                        loaded_history = loaded_histories[prompt]
                        yield cls(
                            id=name,
                            prompt=prompt,
                            reference=gt_reference,
                            env=env,
                            task_type=task_type,
                            loaded_history=loaded_history,
                        )
                    else:
                        LOGGER.info("NO LOADED HISTORY FOUND")
                        continue

        return generator(), N_TASKS

mint/tasks/alfworld.py
xingyaoww/mint-bench
import os
import json
import logging
import functools
from abc import ABC
from typing import List, Optional, Tuple, Mapping

from mint.utils import load_file

LOGGER = logging.getLogger("MINT")


class Task(ABC):
    """Base class for a task instance."""

    task_name: str = "base"
    in_context_example_dir = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "in_context_examples",
    )

    def __init__(self, **kwargs) -> None:
        if "loaded_history" in kwargs:
            self.loaded_history = kwargs["loaded_history"]
        else:
            self.loaded_history = None
        # pre-load the in-context example
        task_dir = os.path.join(self.in_context_example_dir, self.task_name)
        self._in_context_example = {
            "with_tool_and_feedback": load_file(
                os.path.join(task_dir, "with_tool_and_feedback.txt")
            ),
            "with_tool": load_file(os.path.join(task_dir, "with_tool.txt")),
        }
        self.metadata = {}

    @property
    def task_id(self) -> str:
        """Return the task id."""
        assert hasattr(self, "_id"), "Task does not have an id."
        return self._id

    def in_context_example(
        self, use_tool: bool = True, with_feedback: bool = False
    ) -> str:
        """Return the in-context example for the task."""
        if use_tool and with_feedback:
            return self._in_context_example["with_tool_and_feedback"]
        elif use_tool and not with_feedback:
            return self._in_context_example["with_tool"]
        else:
            raise NotImplementedError

    @property
    def prompt(self) -> str:
        """Return the task prompt."""
        assert hasattr(self, "_prompt"), "Task does not have a prompt."
        return self._prompt

    @property
    def reference(self) -> str:
        """Return the reference solution for the task."""
        assert hasattr(self, "_reference"), "Task does not have a reference solution."
        return self._reference

    def extract_answer(self, solution: str) -> Optional[str]:
        """Extract the answer from the given solution."""
        pass

    def success(self, solution: str) -> bool:
        """This checks whether the given solution can complete the current task.

        Can be used to provide binary feedback.
        """
        answer = self.extract_answer(solution)
        return answer == self.reference

    @classmethod
    def load_tasks(cls, path: str) -> Tuple[List["Task"], int]:
        """Load all the tasks from a given jsonl file."""
        assert path.endswith(".jsonl") or path.endswith(".json")
        with open(path, "r") as f:
            tasks = [cls(**json.loads(line)) for line in f.readlines()]
        LOGGER.info(f"Loaded {len(tasks)} tasks from {path}")
        return tasks, len(tasks)

    def to_dict(self) -> dict:
        """Convert the task to a dictionary."""
        return {
            "task_name": self.task_name,
            "task_id": self.task_id,
            "prompt": self.prompt,
            "reference": self.reference,
            "metadata": self.metadata,
        }

mint/tasks/base.py
xingyaoww/mint-bench
import logging
from typing import List, Optional

LOGGER = logging.getLogger("MINT")

from .base import Task
from mint.utils.exec import check_correctness


class CodeGenTask(Task):
    """Generic code generation task instance."""

    def __init__(self, id: str, prompt: str, reference: str, **kwargs):
        super().__init__(**kwargs)
        self._id = id
        self._prompt = prompt
        self._reference = reference

    def success(self, solution: str) -> bool:
        """This checks whether the given solution can complete the current task.

        Can be used to provides binary feedback.
        """
        code_to_exec = self.extract_answer(solution)
        LOGGER.debug(f"CODE_TO_EXEC:\n{code_to_exec}")
        LOGGER.debug(f"TEST_CODE:\n{self._reference}")
        res = check_correctness(
            solution_code=code_to_exec, test_code=self._reference, timeout=10
        )
        return res["success"]


class MBPPTask(CodeGenTask):
    task_name = "mbpp"

    @property
    def prompt(self) -> str:
        """Return the prompt for this task.

        MBPP prompt contains \"\"\" enclosed at both ends. Need to remove it.
        """
        return self._prompt.replace('"""', "").strip()

    def extract_answer(self, solution: str) -> Optional[str]:
        """Extract the answer from the given solution.

        Split off first block of code by scanning for class, def etc. on newlines.

        Modified from:
        https://github.com/bigcode-project/bigcode-evaluation-harness/blob/d61afde130005ecc65cf800ad8eca790a9bc2115/lm_eval/tasks/mbpp.py#L67
        """
        # STOP_WORDS = ["\nclass", "\nassert", '\n"""', "\nprint", "\nif", "\n<|/"]
        # return re.split("|".join(STOP_WORDS), solution)[0].rstrip()
        return solution


class HumanEvalTask(CodeGenTask):

    task_name = "humaneval"

    @property
    def prompt(self) -> str:
        """Return the prompt for this task.

        MBPP prompt contains \"\"\" enclosed at both ends. Need to remove it.
        """
        return "Complete the following code:\n\n" + self._prompt

    def extract_answer(self, solution: str) -> Optional[str]:
        """Extract the answer from the given solution.

        Split off first block of code by scanning for class, def etc. on newlines.

        Modified from:
        https://github.com/bigcode-project/bigcode-evaluation-harness/blob/d61afde130005ecc65cf800ad8eca790a9bc2115/lm_eval/tasks/humaneval.py#L56
        """

        # STOP_WORDS = ["\nclass", "\ndef", "\n#", "\n@", "\nprint", "\nif"]
        # # Remove the last block of the code containing stop_words for HumanEval
        # string_list = re.split("(%s)" % "|".join(STOP_WORDS), solution)
        # # last string should be ""
        # return "".join(string_list[:-2])
        return solution

mint/tasks/codegen.py
xingyaoww/mint-bench
import re
import ast
import logging
import traceback
import numpy as np
from sympy import Rational

from .base import Task
from typing import Optional, Any

LOGGER = logging.getLogger("MINT")


class ReasoningTask(Task):
    task_name = "reasoning"

    def __init__(self, id: str, prompt: str, reference: str, **kwargs):
        super().__init__(**kwargs)
        self._id = id
        self._prompt = prompt.strip()
        self._reference = str(reference).strip().lower()

    def extract_answer(self, solution: str) -> Optional[str]:
        """Extract the answer from the given solution."""
        return solution.lower().strip()

    def compare_w_digits(self, reference: str, answer: str) -> bool:
        """Compare the reference and answer with digits."""
        # if reference can and answer can both be converted to floats by float()
        try:
            float(reference)
            float(answer)
            return abs(float(reference) - float(answer)) <= 0.05 * float(reference)
        except ValueError:
            return reference in answer
        except Exception:
            raise ValueError(f"Cannot compare {reference} and {answer}")

    def success(self, solution: str) -> bool:
        answer = self.extract_answer(solution)
        return self.compare_w_digits(self._reference, answer)


class MultipleChoiceTask(Task):
    """Subclass of Task for multiple choice tasks."""

    task_name = "reasoning"

    def __init__(self, id, prompt: str, reference: str, **kwargs):
        super().__init__(**kwargs)
        self._id = id
        self.hide_options = kwargs.get("hide_options", False)
        if self.hide_options:
            self._prompt = prompt.split("Options:")[0].strip()
        else:
            self._prompt = prompt
        self._reference = reference.strip().lower()
        self._options = self.extract_options(prompt)
        # if all options can be converted to float, strictly perform hide options
        try:
            for option in self._options.values():
                float(option)
            self.hide_options = True
        except ValueError:
            pass
        self.metadata.update({"options": self._options})

    def extract_answer(self, solution: str) -> Optional[str]:
        # Extract the selected option from the solution
        solution = solution.lower().strip()
        for l in "abcdefghijklmnopqrstuvwxyz":
            if f"{l})" in solution or f"{l} )" in solution:
                print("SOLUTION", l)
                return l
            else:
                print("SOLUTION", solution)
                return solution

    def compare_w_digits(self, reference: str, answer: str) -> bool:
        if reference.isdigit() and answer.isdigit():
            return abs(float(reference) - float(answer)) <= 0.05 * float(reference)
        else:
            return reference in answer

    def success(self, solution: str) -> bool:
        answer = self.extract_answer(solution)
        if self.compare_w_digits(self._reference, answer):
            return True
        else:
            correct_option = self._options[self._reference]
            wrong_option_list = list(self._options.values())
            print("OPTIONS", correct_option, wrong_option_list)
            print("ANSWER", answer)
            for i in wrong_option_list:
                if i in correct_option:
                    wrong_option_list.remove(i)
            for i in wrong_option_list:
                if self.compare_w_digits(i, answer) or (i in answer):
                    return False
            if self.compare_w_digits(correct_option, answer) or (
                correct_option in answer
            ):
                return True
            else:
                return False

    def extract_options(self, prompt: str) -> dict:
        # Find the possible option separators (comma, semicolon, or parentheses)
        prompt = prompt.split("Options: ")[-1]
        # Extract the options using the delimiter
        options_match = prompt.split(" , ")
        options = {}
        for i in range(len(options_match)):
            option = options_match[i].strip("[]' ")
            option = option.split(")")
            letter = option[0].lower().strip()
            content = (
                option[1]
                .lower()
                .strip(".")
                .replace(". Which option is correct?", "")
                .replace(". Which one is correct?", "")
                .strip()
            )
            options.update({letter: content})
        return options


# ==== TheoremQA ====


def compare_two_numbers(p, gt):
    if isinstance(p, int) or isinstance(p, float):
        pass
    elif isinstance(p, list) or isinstance(p, bool) or isinstance(p, str):
        return False
    elif isinstance(p, tuple) or isinstance(p, complex) or isinstance(p, dict):
        return False
    else:
        raise ValueError(p)

    if isinstance(gt, float):
        return within_eps(pred=p, gt=gt)
    else:
        return round(p) == gt


def compare_two_list(pred, gt):
    if not isinstance(pred, list):
        return False
    elif len(pred) != len(gt):
        return False
    elif any([not isinstance(x, (int, float)) for x in pred]):
        return False
    else:
        pred = sorted(pred)
        gt = sorted(gt)
        return all([compare_two_numbers(p, g) for p, g in zip(pred, gt)])


def within_eps(pred: float, gt: float):
    eps = abs(gt) * 0.04
    if pred >= gt - eps and pred <= gt + eps:
        return True
    else:
        return False


def parse_number_list(s: str):
    # Check if the string is a valid list by trying to parse it
    parsed_list = ast.literal_eval(s)
    return parsed_list


def is_number(string):
    pattern = r"^[-+]?(\d{1,3}(,\d{3})*|(\d+))(\.\d+)?$"
    match = re.match(pattern, string)
    return bool(match)


def is_scientific_number(string):
    pattern = r"^[-+]?\d+(\.\d+)?e[-]?\d+$"
    match = re.match(pattern, string)
    return bool(match)


def contain_num_and_str(string):
    pattern_str = r"[a-zA-Z]"
    pattern_num = r"[0-9]"
    return bool(re.search(pattern_str, string) and re.search(pattern_num, string))


def is_number(string):
    pattern = r"^[-+]?(\d{1,3}(,\d{3})*|(\d+))(\.\d+)?$"
    match = re.match(pattern, string)
    return bool(match)


def is_scientific_number(string):
    pattern = r"^[-+]?\d+(\.\d+)?e[-]?\d+$"
    match = re.match(pattern, string)
    return bool(match)


class TheoremqaTask(Task):
    task_name = "reasoning"

    def __init__(self, id: str, prompt: str, reference: str, **kwargs):
        super().__init__(**kwargs)
        self._id = id
        self._prompt = (
            "Answer the following question with a number, a list of numbers or True or False. "
            + prompt.strip()
        )
        self._reference = reference
        self._answer_type = kwargs.get("answer_type")

    def extract_answer(self, solution: str) -> Optional[Any]:
        """Extract the answer from the given solution."""
        prediction = solution
        # Following the preprocessing steps from TheoremQA
        # https://github.com/wenhuchen/TheoremQA/blob/123e36beaaa97c01f28a582f13c4f77a6822c199/predict_accuracy.py#L170

        # Preprocessing the string [Stage 1]
        if not isinstance(prediction, str):
            prediction = str(prediction) if prediction is not None else "0"

        # Replace special tokens
        if "=" in prediction:
            prediction = prediction.split("=")[-1].strip()
        if "≈" in prediction:
            prediction = prediction.split("≈")[-1].strip()
        if "`" in prediction:
            prediction = prediction.replace("`", "")
        if "$" in prediction:
            prediction = prediction.replace("$", "")
        if "°" in prediction:
            prediction = prediction.replace("°", "")

        # Detect the boolean keyword in the generation
        if prediction in ["true", "yes", "false", "no"]:
            if prediction == "true" or prediction == "yes":
                prediction = "True"
            else:
                prediction = "False"
        if "True" in prediction or "False" in prediction:
            prediction = "True" if "True" in prediction else "False"

        # Detect the approximation keyword
        if "approximately" in prediction:
            prediction = prediction.replace("approximately", "").strip()
        if " or " in prediction:
            prediction = prediction.split(" or ")[0]

        # Drop the units before and after the number
        if re.match(r"[-+]?(?:[\d,]*\.*\d+) [^0-9 ]+$", prediction):
            prediction = re.search(
                r"([-+]?(?:[\d,]*\.*\d+)) [^0-9 ]+$", prediction
            ).group(1)
        if re.match(r"[^0-9 ]+ [-+]?(?:[\d,]*\.*\d+)$", prediction):
            prediction = re.search(
                r"[^0-9 ]+ ([-+]?(?:[\d,]*\.*\d+))$", prediction
            ).group(1)
        if re.match(r"[-+]?(?:[\d,]*\.*\d+)[^\d]{1,2}$", prediction):
            prediction = re.search(
                r"([-+]?(?:[\d,]*\.*\d+))[^\d]{1,2}$", prediction
            ).group(1)
        if re.match(r"[^-+\d]{1,2}(?:[\d,]*\.*\d+)$", prediction):
            prediction = re.search(
                r"[^-+\d]{1,2}((?:[\d,]*\.*\d+))$", prediction
            ).group(1)

        # Preprocessing the number [Stage 1]
        if "10^" in prediction:
            prediction = re.sub(r"10\^(-?\d+)", r"math.pow(10, \1)", prediction)
        if " x " in prediction:
            prediction = prediction.replace(" x ", "*")
        if " × " in prediction:
            prediction = prediction.replace(" × ", "*")
        if is_number(prediction):
            prediction = prediction.replace(",", "")

        # Preprocessing the option [Stage 3]
        if "a)" in prediction or "a )" in prediction or prediction.lower().strip() == "a":
            prediction = "(a)"
        if "b)" in prediction or "b )" in prediction or prediction.lower().strip() == "b":
            prediction = "(b)"
        if "c)" in prediction or "c )" in prediction or prediction.lower().strip() == "c":
            prediction = "(c)"
        if "d)" in prediction or "d )" in prediction or prediction.lower().strip() == "d":
            prediction = "(d)"

        if (
            "(a)" in prediction
            or "(b)" in prediction
            or "(c)" in prediction
            or "(d)" in prediction
        ):
            prediction = '"' + re.search(r"\([a-d]\)", prediction).group(0) + '"'

        # If the prediction is empty, use dummy '0'
        if not prediction:
            prediction = "0"

        # Converting the string answer to a number/list/bool/option
        try:
            prediction = eval(prediction)
        except Exception as e:
            LOGGER.warning(
                f"[TASK] Failed to convert the answer: {prediction}\n{traceback.format_exc()}"
            )
            return None  # failed to convert the answer

        # Performing common type conversion
        if isinstance(prediction, (set, tuple)):
            prediction = list(prediction)
            if isinstance(prediction[0], complex):
                prediction = [tmp.real for tmp in prediction]
            elif isinstance(prediction[0], Rational):
                prediction = [float(tmp) for tmp in prediction]
        elif isinstance(prediction, np.ndarray):
            prediction = prediction.tolist()
        else:
            if isinstance(prediction, complex):
                prediction = prediction.real
            elif isinstance(prediction, Rational):
                prediction = float(prediction)

        return prediction

    def success(self, solution: str) -> bool:
        """This checks whether the given solution can complete the current task."""
        # Follow the implementation from TheoremQA
        # https://github.com/wenhuchen/TheoremQA/blob/123e36beaaa97c01f28a582f13c4f77a6822c199/predict_accuracy.py#L301C9-L317C1
        prediction = self.extract_answer(solution)
        LOGGER.info(f"TheoremQA Parsed Prediction: {prediction}")
        answer_type = self._answer_type
        gt = self.reference

        if isinstance(prediction, (str, int, float)) or isinstance(prediction, list):
            # Comparing prediction against the reference
            if answer_type in ["bool", "option", "Option"]:
                cur_correct = int(prediction == f"({gt})") or int(prediction == gt)
            elif answer_type == "integer":
                cur_correct = int(compare_two_numbers(prediction, gt))
            elif answer_type == "float":
                cur_correct = int(compare_two_numbers(prediction, gt))
            elif answer_type in ["list of integer", "list of float"]:
                cur_correct = int(compare_two_list(prediction, gt))
        else:
            cur_correct = 0
        return bool(cur_correct)

mint/tasks/reasoning.py
xingyaoww/mint-bench
from .base import Tool
from typing import List
import warnings

warnings.filterwarnings("ignore")


def get_toolset_description(tools: List[Tool]) -> str:
    if len(tools) == 0:
        return ""

    output = "Tool function available (already imported in <execute> environment):\n"
    for i, tool in enumerate(tools):
        output += f"[{i + 1}] {tool.signature}\n"
        output += f"{tool.description}\n"

    return output

mint/tools/__init__.py
xingyaoww/mint-bench
from mint.tools.base import Tool
from mint.datatypes import StepOutput


def process_ob(ob):
    if ob.startswith("You arrive at loc "):
        ob = ob[ob.find(". ") + 2 :]
    return ob


class AlfWorldTool(Tool):
    """Base class for a tool in AlfWorld."""

    name = ""
    signature = ""
    description = ""

    def __init__(self, env, callback_fn: callable) -> None:
        self.env = env
        self.callback_fn = callback_fn

    def __call__(self, **kwargs) -> StepOutput:
        pass

    def step(self, action: str) -> str:
        # Follow ReAct Implementation
        observation, reward, done, info = self.env.step([action])
        observation, reward, done = process_ob(observation[0]), info["won"][0], done[0]
        self.callback_fn(
            StepOutput(
                observation=observation,
                success=done,
                extra={"reward": reward},
            )
        )
        # print(f"After '{action}': {observation}")
        return observation


SIGNATURE_TEMPLATE = "{name}({args}) -> str"

DESCRIPTION_TEMPLATE = """{desc}
For example: {example}
""".lstrip()


class Put(AlfWorldTool):
    name = "put"
    signature = SIGNATURE_TEMPLATE.format(
        name=name, args="object: str, receptacle: str"
    )
    description = DESCRIPTION_TEMPLATE.format(
        desc="Put an object in/on a receptacle.",
        example='put("mug 1", "desk 2")',
    )

    def __call__(self, object: str, receptacle: str):
        return self.step(f"put {object} in/on {receptacle}")


class GoTo(AlfWorldTool):
    name = "goto"
    signature = SIGNATURE_TEMPLATE.format(name=name, args="receptacle: str")
    description = DESCRIPTION_TEMPLATE.format(
        desc="Go to a location of the receptacle.",
        example='goto("drawer 1")',
    )

    def __call__(self, receptacle: str):
        return self.step(f"go to {receptacle}")


class Take(AlfWorldTool):
    name = "take_from"
    signature = SIGNATURE_TEMPLATE.format(
        name=name, args="object: str, receptacle: str"
    )
    description = DESCRIPTION_TEMPLATE.format(
        desc="Take an object from a receptacle.",
        example='take_from("mug 1", "shelf 2")',
    )

    def __call__(self, object: str, receptacle: str):
        return self.step(f"take {object} from {receptacle}")


class Open(AlfWorldTool):
    name = "open_receptacle"
    signature = SIGNATURE_TEMPLATE.format(name=name, args="receptacle: str")
    description = DESCRIPTION_TEMPLATE.format(
        desc="Open a receptacle.",
        example='open_receptacle("fridge 1")',
    )

    def __call__(self, receptacle: str):
        return self.step(f"open {receptacle}")


class Toggle(AlfWorldTool):
    name = "toggle"
    signature = SIGNATURE_TEMPLATE.format(name=name, args="object_or_receptacle: str")
    description = DESCRIPTION_TEMPLATE.format(
        desc="Toggle an object or receptacle.",
        example='toggle("light 2")',
    )

    def __call__(self, object_or_receptacle: str):
        return self.step(f"toggle {object_or_receptacle}")


class Close(AlfWorldTool):
    name = "close_receptacle"
    signature = SIGNATURE_TEMPLATE.format(
        name=name,
        args="receptacle: str",
    )
    description = DESCRIPTION_TEMPLATE.format(
        desc="Close a receptacle.",
        example='close_receptacle("microwave 1")',
    )

    def __call__(self, receptacle: str):
        return self.step(f"close {receptacle}")


class Clean(AlfWorldTool):
    name = "clean"
    signature = SIGNATURE_TEMPLATE.format(
        name=name,
        args="object: str, receptacle: str",
    )
    description = DESCRIPTION_TEMPLATE.format(
        desc="Clean an object with a receptacle.",
        example='clean("cloth 1", "sinkbasin 1")',
    )

    def __call__(self, object: str, receptacle: str):
        return self.step(f"clean {object} with {receptacle}")


class Heat(AlfWorldTool):
    name = "heat"
    signature = SIGNATURE_TEMPLATE.format(
        name=name,
        args="object: str, receptacle: str",
    )
    description = DESCRIPTION_TEMPLATE.format(
        desc="Heat an object with a receptacle.",
        example='heat("egg 1", "microwave 1")',
    )

    def __call__(self, object: str, receptacle: str):
        return self.step(f"heat {object} with {receptacle}")


class Cool(AlfWorldTool):
    name = "cool"
    signature = SIGNATURE_TEMPLATE.format(
        name=name,
        args="object: str, receptacle: str",
    )
    description = DESCRIPTION_TEMPLATE.format(
        desc="Cool an object with a receptacle.",
        example='cool("bottle 1", "fridge 1")',
    )

    def __call__(self, object: str, receptacle: str):
        return self.step(f"cool {object} with {receptacle}")


# add use tool
class Use(AlfWorldTool):
    name = "use"
    signature = SIGNATURE_TEMPLATE.format(
        name=name,
        args="receptacle: str",
    )
    description = DESCRIPTION_TEMPLATE.format(
        desc="Use a receptacle.",
        example='use("lamp 1")',
    )

    def __call__(self, receptacle: str):
        return self.step(f"use {receptacle}")


# add look tool
class Look(AlfWorldTool):
    name = "look"
    signature = SIGNATURE_TEMPLATE.format(
        name=name,
        args="",
    )
    description = DESCRIPTION_TEMPLATE.format(
        desc="Look around. It will return what you see in the room.",
        example="look()",
    )

    def __call__(self):
        return self.step(f"look")


ALFWORLD_TOOL_CLS = [
    Put,
    GoTo,
    Take,
    Open,
    Toggle,
    Close,
    Clean,
    Heat,
    Cool,
    Use,
    Look,
]

mint/tools/alfworld.py
xingyaoww/mint-bench
from abc import ABC, abstractmethod
from typing import Any


class Tool(ABC):
    """Abstract class for a tool."""

    name: str
    signature: str
    description: str

    @abstractmethod
    def __call__(self, *args: Any, **kwds: Any) -> str:
        """Execute the tool with the given args and return the output."""
        # execute tool with abitrary args
        pass

    def reset(self) -> None:
        """Reset the tool to its initial state."""
        pass

mint/tools/base.py
xingyaoww/mint-bench
from typing import Mapping
from mint.tools.base import Tool
import re
import signal
from contextlib import contextmanager
from IPython.core.interactiveshell import InteractiveShell
from IPython.utils import io


class PythonREPL(Tool):
    """A tool for running python code in a REPL."""

    name = "PythonREPL"
    # This PythonREPL is not used by the environment; It is THE ENVIRONMENT.
    signature = "NOT_USED"
    description = "NOT_USED"

    def __init__(
        self,
        name_to_func_mapping: Mapping[str, callable],
        timeout: int = 30,
    ) -> None:
        super().__init__()
        self.user_ns = name_to_func_mapping
        self.timeout = timeout
        self.reset()

    @contextmanager
    def time_limit(self, seconds):
        def signal_handler(signum, frame):
            raise TimeoutError(f"Timed out after {seconds} seconds.")

        signal.signal(signal.SIGALRM, signal_handler)
        signal.alarm(seconds)
        try:
            yield
        finally:
            signal.alarm(0)  # Disable the alarm

    def reset(self) -> None:
        InteractiveShell.clear_instance()
        self.shell = InteractiveShell.instance(
            # NOTE: shallow copy is needed to avoid
            # shell modifying the original user_ns dict
            user_ns=dict(self.user_ns),
            colors="NoColor",
        )

    def __call__(self, query: str) -> str:
        """Use the tool and return observation"""
        with self.time_limit(self.timeout):
            # NOTE: The timeout error will be caught by the InteractiveShell

            # Capture all output
            with io.capture_output() as captured:
                _ = self.shell.run_cell(query, store_history=True)
            output = captured.stdout

            if output == "":
                output = "[Executed Successfully with No Output]"

            # replace potentially sensitive filepath
            # e.g., File /mint/mint/tools/python_tool.py:30, in PythonREPL.time_limit.<locals>.signal_handler(signum, frame)
            # with File <filepath>:30, in PythonREPL.time_limit.<locals>.signal_handler(signum, frame)
            # use re
            output = re.sub(
                # r"File (/mint/)mint/tools/python_tool.py:(\d+)",
                r"File (.*)mint/tools/python_tool.py:(\d+)",
                r"File <hidden_filepath>:\1",
                output,
            )
            if len(output) > 2000:
                output = output[:2000] + "...\n[Output Truncated]"

        return output


if __name__ == "__main__":
    shell = PythonREPL(
        name_to_func_mapping={"hello": lambda: print("hello world")},
        timeout=5,
    )

    def exec_tool(input_code: str, cur_shell=shell) -> str:
        print("=== INPUT ===")
        print(input_code)
        print("=== OUTPUT ===")
        output = cur_shell(input_code)
        print(output)

    INPUT = """
import time
i = 0
while True:
    time.sleep(1)
    i += 1
    print(i)
"""
    exec_tool(INPUT)

    INPUT = """
a = 1
b = 2
a + b
"""
    exec_tool(INPUT)

    # test whether variables `a` persists
    INPUT = """a + 7"""
    exec_tool(INPUT)

    # test provided import
    INPUT = """hello()"""
    exec_tool(INPUT)

    print("\n=== SECOND SHELL ===\n")
    second_shell = PythonREPL(
        name_to_func_mapping={"hello": lambda: print("not hello world!")},
        timeout=5,
    )

    INPUT = """a + 7"""
    exec_tool(INPUT, cur_shell=second_shell)

    INPUT = """hello()"""
    exec_tool(INPUT, cur_shell=second_shell)

    # reset tool and test whether variables `a` persists
    shell.reset()

    print("\n=== 1st Shell RESET ===\n")

    INPUT = """a + 7"""
    exec_tool(INPUT)

    # test package import
    INPUT = """
import math
math.sqrt(4)
"""
    exec_tool(INPUT)

    # test provided import after reset
    INPUT = """hello()"""
    exec_tool(INPUT)

    # code that leads to traceback
    INPUT = """
def similar_elements(test_tup1, test_tup2):
    res = tuple(set(test_tup1) | set(test_tup2))
    return res
res = similar_elements((3, 4, 5, 6), (5, 7, 4, 10))
assert res == (4, 5), "Expected (4, 5) but got {}".format(res)
"""
    exec_tool(INPUT)

    # code with no error
    INPUT = """
def similar_elements(test_tup1, test_tup2):
    res = tuple(set(test_tup1) & set(test_tup2))
    return res
res = similar_elements((3, 4, 5, 6), (5, 7, 4, 10))
assert res == (4, 5), "Expected (4, 5) but got {}".format(res)
"""
    exec_tool(INPUT)

mint/tools/python_tool.py
xingyaoww/mint-bench
"""Util that calls Wikipedia."""
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Extra, root_validator
from mint.tools.base import Tool
from langchain.schema import Document

WIKIPEDIA_MAX_QUERY_LENGTH = 300

class WikipediaAPIWrapper(BaseModel):
    """Wrapper around WikipediaAPI.

    To use, you should have the ``wikipedia`` python package installed.
    This wrapper will use the Wikipedia API to conduct searches and
    fetch page summaries. By default, it will return the page summaries
    of the top-k results.
    It limits the Document content by doc_contenrt_chars_max.
    """

    wiki_client: Any  #: :meta private:
    top_k_results: int = 2
    lang: str = "en"
    load_all_available_meta: bool = False
    doc_content_chars_max: int = 1000

    class Config:
        """Configuration for this pydantic object."""

        extra = Extra.forbid

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """Validate that the python package exists in environment."""
        try:
            import wikipedia

            wikipedia.set_lang(values["lang"])
            values["wiki_client"] = wikipedia
        except ImportError:
            raise ImportError(
                "Could not import wikipedia python package. "
                "Please install it with `pip install wikipedia`."
            )
        return values

    def run(self, query: str) -> str:
        """Run Wikipedia search and get page summaries."""
        page_titles = self.wiki_client.search(query[:WIKIPEDIA_MAX_QUERY_LENGTH])
        summaries = []
        for page_title in page_titles[: self.top_k_results]:
            if wiki_page := self._fetch_page(page_title):
                if summary := self._formatted_page_summary(page_title, wiki_page):
                    summaries.append(summary)
        if not summaries:
            return "Wikipedia Search Result was found"
        return "\n\n".join(summaries)[: self.doc_content_chars_max]

    @staticmethod
    def _formatted_page_summary(page_title: str, wiki_page: Any) -> Optional[str]:
        return f"Page: {page_title}\nSummary: {wiki_page.summary}"

    def _page_to_document(self, page_title: str, wiki_page: Any) -> Document:
        main_meta = {
            "title": page_title,
            "summary": wiki_page.summary,
            "source": wiki_page.url,
        }
        add_meta = (
            {
                "categories": wiki_page.categories,
                "page_url": wiki_page.url,
                "image_urls": wiki_page.images,
                "related_titles": wiki_page.links,
                "parent_id": wiki_page.parent_id,
                "references": wiki_page.references,
                "revision_id": wiki_page.revision_id,
                "sections": wiki_page.sections,
            }
            if self.load_all_available_meta
            else {}
        )
        doc = Document(
            page_content=wiki_page.content[: self.doc_content_chars_max],
            metadata={
                **main_meta,
                **add_meta,
            },
        )
        return doc

    def _fetch_page(self, page: str) -> Optional[str]:
        try:
            return self.wiki_client.page(title=page, auto_suggest=False)
        except (
            self.wiki_client.exceptions.PageError,
            self.wiki_client.exceptions.DisambiguationError,
        ):
            return None

    def load(self, query: str) -> List[Document]:
        """
        Run Wikipedia search and get the article text plus the meta information.
        See

        Returns: a list of documents.

        """
        page_titles = self.wiki_client.search(query[:WIKIPEDIA_MAX_QUERY_LENGTH])
        docs = []
        for page_title in page_titles[: self.top_k_results]:
            if wiki_page := self._fetch_page(page_title):
                if doc := self._page_to_document(page_title, wiki_page):
                    docs.append(doc)
        return docs


WIKIPEDIA_DESCRIPTION = """The Wikipedia Search tool provides access to a vast collection of articles covering a wide range of topics.
Can query specific keywords or topics to retrieve accurate and comprehensive information.
"""


class WikipediaQueryRun(Tool):
    """Tool that adds the capability to search using the Wikipedia API."""

    name = "wikipedia_search"
    signature = f"{name}(query: str) -> str"
    description = WIKIPEDIA_DESCRIPTION
    api_wrapper = WikipediaAPIWrapper()

    def __call__(
        self,
        query: str,
    ) -> str:
        """Use the Wikipedia tool."""
        output = self.api_wrapper.run(query)
        return output

mint/tools/wikipedia_search.py
xingyaoww/mint-bench
import functools

# use cache to avoid loading the same file multiple times
# which can leads to too many open files error
@functools.lru_cache(maxsize=128)
def load_file(filepath: str) -> str:
    with open(filepath, "r") as f:
        content = f.read()
    return content

mint/utils/__init__.py
xingyaoww/mint-bench
class ParseError(Exception):
    pass


class ToolExecutionError(Exception):
    pass


class LMExecutionTimeoutError(Exception):
    pass

mint/utils/exception.py
xingyaoww/mint-bench
"""Check the correctness of a program by running a test suite.

Modified from: https://github.com/openai/human-eval/blob/master/human_eval/execution.py
"""

from typing import Optional, Callable, Dict
import ast
import contextlib
import faulthandler
import io
import os
import multiprocessing
import platform
import signal
import tempfile


def check_correctness(
    solution_code: str,
    test_code: str,
    timeout: float = 10,
    completion_id: Optional[int] = None,
) -> Dict:
    """
    Evaluates the functional correctness of a completion by running the test
    suite provided in the problem.

    :param completion_id: an optional completion ID so we can match
        the results later even if execution finishes asynchronously.
    """

    def unsafe_execute():
        with create_tempdir():
            # These system calls are needed when cleaning up tempdir.
            import os
            import shutil

            rmtree = shutil.rmtree
            rmdir = os.rmdir
            chdir = os.chdir

            # Disable functionalities that can make destructive changes to the test.
            reliability_guard()

            # Construct the check program and run it.
            check_program = solution_code + "\n" + test_code

            try:
                exec_globals = {}
                with swallow_io():
                    with time_limit(timeout):
                        # WARNING
                        # This program exists to execute untrusted model-generated code. Although
                        # it is highly unlikely that model-generated code will do something overtly
                        # malicious in response to this test suite, model-generated code may act
                        # destructively due to a lack of model capability or alignment.
                        # Users are strongly encouraged to sandbox this evaluation suite so that it
                        # does not perform destructive actions on their host or network. For more
                        # information on how OpenAI sandboxes its code, see the accompanying paper.
                        # Once you have read this disclaimer and taken appropriate precautions,
                        # uncomment the following line and proceed at your own risk:
                        exec(check_program, exec_globals)
                result.append("passed")
            except TimeoutException:
                result.append("timed out")
            except BaseException as e:
                result.append(f"failed: {e}")

            # Needed for cleaning up.
            shutil.rmtree = rmtree
            os.rmdir = rmdir
            os.chdir = chdir

    manager = multiprocessing.Manager()
    result = manager.list()

    p = multiprocessing.Process(target=unsafe_execute)
    p.start()
    p.join(timeout=timeout + 1)
    if p.is_alive():
        p.kill()

    if not result:
        result.append("timed out")

    return dict(
        success=result[0] == "passed",
        result=result[0],
        completion_id=completion_id,
    )


@contextlib.contextmanager
def time_limit(seconds: float):
    def signal_handler(signum, frame):
        raise TimeoutException("Timed out!")

    signal.setitimer(signal.ITIMER_REAL, seconds)
    signal.signal(signal.SIGALRM, signal_handler)
    try:
        yield
    finally:
        signal.setitimer(signal.ITIMER_REAL, 0)


@contextlib.contextmanager
def swallow_io():
    stream = WriteOnlyStringIO()
    with contextlib.redirect_stdout(stream):
        with contextlib.redirect_stderr(stream):
            with redirect_stdin(stream):
                yield


@contextlib.contextmanager
def create_tempdir():
    # with tempfile.TemporaryDirectory() as dirname:
    # Manually do this to avoid too many open files error caused by TemporaryDirectory
    dirname = tempfile.mkdtemp()
    with chdir(dirname):
        yield dirname
    os.rmdir(dirname)


class TimeoutException(Exception):
    pass


class WriteOnlyStringIO(io.StringIO):
    """StringIO that throws an exception when it's read from"""

    def read(self, *args, **kwargs):
        raise IOError

    def readline(self, *args, **kwargs):
        raise IOError

    def readlines(self, *args, **kwargs):
        raise IOError

    def readable(self, *args, **kwargs):
        """Returns True if the IO object can be read."""
        return False


class redirect_stdin(contextlib._RedirectStream):  # type: ignore
    _stream = "stdin"


@contextlib.contextmanager
def chdir(root):
    if root == ".":
        yield
        return
    cwd = os.getcwd()
    os.chdir(root)
    try:
        yield
    except BaseException as exc:
        raise exc
    finally:
        os.chdir(cwd)


def reliability_guard(maximum_memory_bytes: Optional[int] = None):
    """
    This disables various destructive functions and prevents the generated code
    from interfering with the test (e.g. fork bomb, killing other processes,
    removing filesystem files, etc.)

    WARNING
    This function is NOT a security sandbox. Untrusted code, including, model-
    generated code, should not be blindly executed outside of one. See the
    Codex paper for more information about OpenAI's code sandbox, and proceed
    with caution.
    """

    if maximum_memory_bytes is not None:
        import resource

        resource.setrlimit(
            resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes)
        )
        resource.setrlimit(
            resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes)
        )
        if not platform.uname().system == "Darwin":
            resource.setrlimit(
                resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes)
            )

    faulthandler.disable()

    import builtins

    builtins.exit = None
    builtins.quit = None

    import os

    os.environ["OMP_NUM_THREADS"] = "1"

    os.kill = None
    os.system = None
    os.putenv = None
    os.remove = None
    os.removedirs = None
    os.rmdir = None
    os.fchdir = None
    os.setuid = None
    os.fork = None
    os.forkpty = None
    os.killpg = None
    os.rename = None
    os.renames = None
    os.truncate = None
    os.replace = None
    os.unlink = None
    os.fchmod = None
    os.fchown = None
    os.chmod = None
    os.chown = None
    os.chroot = None
    os.fchdir = None
    os.lchflags = None
    os.lchmod = None
    os.lchown = None
    os.getcwd = None
    os.chdir = None

    import shutil

    shutil.rmtree = None
    shutil.move = None
    shutil.chown = None

    import subprocess

    subprocess.Popen = None  # type: ignore

    __builtins__["help"] = None

    import sys

    sys.modules["ipdb"] = None
    sys.modules["joblib"] = None
    sys.modules["resource"] = None
    sys.modules["psutil"] = None
    sys.modules["tkinter"] = None

mint/utils/exec.py
xingyaoww/mint-bench
from glob import glob
files = glob("nas/mint-bench/data/outputs/**/theoremqa/results.jsonl", recursive=True)
import json
for file in files:
    json_strs = open(file).readlines()
    good_jsons = []
    for json_str in json_strs:
        json_obj = json.loads(json_str)
        possible_solutions = [i['content'].split('<solution>')[-1].split('</solution>')[0].strip() for i in json_obj['state']['history']]
        possible_solutions = [i for i in possible_solutions if i in ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j']]
        if len(possible_solutions) == 0:
            good_jsons.append(json_str)
    # dump back
    with open(file, 'w') as f:
        for json_str in good_jsons:
            f.write(json_str)

scripts/clean_outputs.py
xingyaoww/mint-bench
"""Streamlit visualizer.

streamlit run scripts/visualizer/visualizer.py -- --data_dir data/outputs
"""
import re
import json
import random
import pandas as pd
import streamlit as st
import altair as alt
import matplotlib.pyplot as plt
from glob import glob
from pandas.api.types import (
    is_categorical_dtype,
    is_datetime64_any_dtype,
    is_numeric_dtype,
    is_object_dtype,
)
import pandas as pd
import streamlit as st

# default wide mode
st.set_page_config(layout="wide", page_title="Interaction-Eval Output Visualizer")

st.title("Interaction-Eval Output Visualizer")

# Select your data directory
query_params = st.experimental_get_query_params()
glob_pattern = st.text_input(
    "Enter the glob to your data directory.\nFor example: `/path/to/your/mint/data/outputs/CodeLlama*/F=None/max5_p2+tool+cd/**/*results.jsonl`",
    key="data_dir",
    value=query_params.get(
        "data_dir",
        ["/path/to/your/mint/data/outputs/**/*results.jsonl"],
    )[0],
)
st.experimental_set_query_params(glob_pattern=glob_pattern)

filepaths = list(set(glob(glob_pattern, recursive=True)))
st.write(f"Matching glob pattern: `{glob_pattern}`. **{len(filepaths)}** files found.")


def parse_filepath(filepath):
    # gpt-3.5-turbo-0613/F=gpt-3.5-turbo-16k-0613/PHF=no_GT-textual/max5_p2+tool+cd/code_generation/humaneval/results.jsonl
    # gpt-3.5-turbo-0613/F=None/max5_p2+tool+cd/code_generation/humaneval/results.jsonl
    splited = (
        filepath[filepath.find("data/outputs/") + len("data/outputs/") :]
        .lstrip("/")
        .split("/")
    )
    try:
        agent_model_name = splited[0]
        feedback_model_name = splited[1].split("=")[1]
        if feedback_model_name != "None":
            feedback_setting = splited[2]
        else:
            feedback_setting = "None"
        task_name = splited[-2]
        task_type = splited[-3]
        exp_setting = splited[-4]
        return {
            "agent_model_name": agent_model_name,
            "feedback_model_name": feedback_model_name,
            "feedback_setting": feedback_setting,
            "task_name": task_name,
            "task_type": task_type,
            "exp_setting": exp_setting,
            "filepath": filepath,
        }
    except Exception as e:
        st.write([filepath, e, splited])


def filter_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    """
    Adds a UI on top of a dataframe to let viewers filter columns

    Args:
        df (pd.DataFrame): Original dataframe

    Returns:
        pd.DataFrame: Filtered dataframe
    """
    modify = st.checkbox("Add filters")

    if not modify:
        return df

    df = df.copy()

    # Try to convert datetimes into a standard format (datetime, no timezone)
    for col in df.columns:
        if is_object_dtype(df[col]):
            try:
                df[col] = pd.to_datetime(df[col])
            except Exception:
                pass

        if is_datetime64_any_dtype(df[col]):
            df[col] = df[col].dt.tz_localize(None)

    modification_container = st.container()

    with modification_container:
        to_filter_columns = st.multiselect("Filter dataframe on", df.columns)
        for column in to_filter_columns:
            left, right = st.columns((1, 20))
            # Treat columns with < 10 unique values as categorical
            if is_categorical_dtype(df[column]) or df[column].nunique() < 10:
                user_cat_input = right.multiselect(
                    f"Values for {column}",
                    df[column].unique(),
                    default=list(df[column].unique()),
                )
                df = df[df[column].isin(user_cat_input)]
            elif is_numeric_dtype(df[column]):
                _min = float(df[column].min())
                _max = float(df[column].max())
                step = (_max - _min) / 100
                user_num_input = right.slider(
                    f"Values for {column}",
                    min_value=_min,
                    max_value=_max,
                    value=(_min, _max),
                    step=step,
                )
                df = df[df[column].between(*user_num_input)]
            elif is_datetime64_any_dtype(df[column]):
                user_date_input = right.date_input(
                    f"Values for {column}",
                    value=(
                        df[column].min(),
                        df[column].max(),
                    ),
                )
                if len(user_date_input) == 2:
                    user_date_input = tuple(map(pd.to_datetime, user_date_input))
                    start_date, end_date = user_date_input
                    df = df.loc[df[column].between(start_date, end_date)]
            else:
                user_text_input = right.text_input(
                    f"Substring or regex in {column}",
                )
                if user_text_input:
                    df = df[df[column].astype(str).str.contains(user_text_input)]

    return df


def dataframe_with_selections(
    df,
    selected_values=None,
    selected_col="filepath",
):
    # https://docs.streamlit.io/knowledge-base/using-streamlit/how-to-get-row-selections
    df_with_selections = df.copy()
    df_with_selections.insert(0, "Select", False)

    # Set the initial state of "Select" column based on query parameters
    if selected_values:
        df_with_selections.loc[
            df_with_selections[selected_col].isin(selected_values), "Select"
        ] = True

    # Get dataframe row-selections from user with st.data_editor
    edited_df = st.data_editor(
        df_with_selections,
        hide_index=True,
        column_config={"Select": st.column_config.CheckboxColumn(required=True)},
        disabled=df.columns,
    )

    # Filter the dataframe using the temporary column, then drop the column
    selected_rows = edited_df[edited_df.Select]
    return selected_rows.drop("Select", axis=1)


filepaths = pd.DataFrame(list(map(parse_filepath, filepaths)))

# ===== Select a file to visualize =====

filepaths = filepaths.sort_values(
    [
        "agent_model_name",
        "feedback_model_name",
        "task_type",
        "task_name",
        "feedback_setting",
        "exp_setting",
    ]
)

st.markdown("**Select file(s) to visualize**")
filepaths = filter_dataframe(filepaths)
# Make these two buttons are on the same row
# col1, col2 = st.columns(2)
col1, col2 = st.columns([0.08, 1])
select_all = col1.button("Select all")
deselect_all = col2.button("Deselect all")
selected_values = query_params.get("filepaths", None)
selected_values = filepaths["filepath"].tolist() if select_all else selected_values
selected_values = [] if deselect_all else selected_values

selection = dataframe_with_selections(
    filepaths,
    selected_values=selected_values,
    selected_col="filepath",
)
# st.write("Your selection:")
# st.write(selection)
select_filepaths = selection["filepath"].tolist()
# update query params
st.experimental_set_query_params(glob_pattern=glob_pattern, filepaths=select_filepaths)

data = []
for filepath in select_filepaths:
    with open(filepath, "r") as f:
        for l in f.readlines():
            data.append(json.loads(l))
st.write(f"{len(data)} rows found.")

# ===== Task-level dashboard =====


def agg_stats(data):
    stats = []
    for idx, entry in enumerate(data):
        # if len(entry["state"]["history"]) % 2 != 0: continue
        task = {
            k: v for k, v in entry["task"].items() if k not in ["prompt", "reference"]
        }
        if "metadata" in task:
            for k, v in task["metadata"].items():
                task[k] = v
            del task["metadata"]

        stats.append(
            {
                "idx": idx,
                "success": entry["state"]["success"],
                "n_turns": len(entry["state"]["history"]) // 2,
                "terminate_reason": entry["state"]["terminate_reason"],
                "agent_action_count": entry["state"]["agent_action_count"],
                **task,
            }
        )
    return pd.DataFrame(stats)


st.markdown("---")
stats_df = agg_stats(data)
if len(stats_df) == 0:
    st.write("No data to visualize.")
    st.stop()
success_count = stats_df["success"].sum()
st.markdown(
    f"**Success Rate: {success_count / len(data):2%}**: {success_count} / {len(data)} rows are successful."
)


def plot_stats(stats_df):
    # 1. Plot a distribution of terminate_reason, also show percentage
    # fig, ax = plt.subplots(figsize=(6, 1))
    # sns.countplot(y="terminate_reason", data=stats_df, ax=ax)
    # ax.set_title("Distribution of Terminate Reason")
    # ax.set_xlabel("Count")
    # ax.set_ylabel("")
    # # add percentage
    # total = len(stats_df)
    # for p in ax.patches:
    #     percentage = "{:.1f}%".format(100 * p.get_width() / total)
    #     x = p.get_x() + p.get_width()
    #     y = p.get_y() + p.get_height() / 2
    #     ax.annotate(percentage, (x, y))
    # st.pyplot(fig, use_container_width=False)
    terminate_reason_count = (
        stats_df["terminate_reason"]
        .value_counts()
        .reset_index()
        .rename(columns={"index": "terminate_reason", "terminate_reason": "count"})
    )
    chart = (
        alt.Chart(terminate_reason_count)
        .mark_bar()
        .encode(
            x=alt.X("count", type="quantitative", title=""),
            y=alt.Y("terminate_reason", type="nominal", title=""),
            # color=alt.Color("variable", type="nominal", title=""),
            # order=alt.Order("variable", sort="descending"),
        )
    )
    st.altair_chart(chart, use_container_width=True)

    # 1. plot success rate changes over n_turns
    fig, ax = plt.subplots(figsize=(6, 3))
    max_turns = stats_df["n_turns"].max()
    _df = []
    for turn_id in range(1, max_turns + 1):
        _df.append(
            {
                "n_turns": turn_id,
                "success_rate": stats_df[stats_df["n_turns"] <= turn_id][
                    "success"
                ].sum()
                / len(stats_df),
            }
        )
    _df = pd.DataFrame(_df)
    # sns.lineplot(x="n_turns", y="success_rate", data=_df, ax=ax)
    # ax.set_title("Task Success Rate vs. Number of Turns")
    # ax.set_xlabel("Number of Turns")
    # ax.set_ylabel("Success Rate")
    # make xlabel integer
    # ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True))
    # st.pyplot(fig, use_container_width=False)

    chart = (
        alt.Chart(_df, title="Task Success Rate vs. Number of Turns")
        .mark_line()
        .encode(
            x=alt.X("n_turns", type="quantitative", title="Number of Turns"),
            y=alt.Y("success_rate", type="quantitative", title="Success Rate"),
        )
        .properties(width=550)
    )
    st.altair_chart(chart, use_container_width=True)


with st.expander("See figures", expanded=True):
    plot_stats(stats_df)

# ===== Select a row to visualize =====
st.markdown("---")
with st.expander("Stats Table", expanded=True):

    # Add a button to randomly select a row
    if st.button("Randomly Select a Row"):
        row_id = random.choice(stats_df["idx"].values)
        query_params["row_idx"] = [row_id]

    selected_row = dataframe_with_selections(
        stats_df,
        query_params.get("row_idx", None),
        selected_col="idx",
    )
    if len(selected_row) == 0:
        st.write("No row selected.")
        st.stop()
    elif len(selected_row) > 1:
        st.write("More than one row selected.")
        st.stop()
    row_id = selected_row["idx"].values[0]

    # update query params
    st.experimental_set_query_params(
        glob_pattern=glob_pattern, filepaths=select_filepaths, row_idx=[row_id]
    )

# row_id = st.number_input(
#     "Select a row to visualize", min_value=0, max_value=len(data) - 1, value=0
# )
# # row = df.iloc[row_id]

# ===== Visualize the row =====
st.write(f"Visualizing row `{row_id}`")
row_dict = data[row_id]

n_turns = len(row_dict["state"]["history"]) // 2
st.write(f"Number of turns: {n_turns}")


def visualize_row(row_dict):
    # 1. Visualize the task and metrics
    st.markdown("## Task")
    task_dict = row_dict["task"]
    st.markdown("**Prompt**")
    st.write(task_dict["prompt"])
    st.markdown("**Reference**")
    st.write(task_dict["reference"])
    st.json({k: v for k, v in task_dict.items() if k not in ["prompt", "reference"]})

    # 2. Visualize the state
    st.markdown("## State")
    state_dict = row_dict["state"]
    state_dict.pop("latest_output")
    history = state_dict.pop("history")
    st.json(state_dict)
    st.markdown("### History")

    raw_history = "\n".join(
        [f"=== {turn['role']} ===\n{turn['content']}" for turn in history]
    )
    expandable = st.expander("Show raw history", expanded=False)
    with expandable:
        st.code(raw_history)

    ROLES = [
        ("system", "blue"),
        ("user", "red"),
        ("assistant", "green"),
    ]
    REGEX_TO_HIGHLIGHT = [
        (r"<thought>(.*?)</thought>", "blue", "thought"),
        (r"<solution>(.*?)</solution>", "green", "solution"),
        (r"<execute>(.*?)</execute>", "orange", "execute"),
    ]
    for turn_id, turn_content in enumerate(history):
        if turn_id % 2 == 0:
            st.markdown(f"#### Turn {(turn_id // 2) + 1}")

        role = turn_content["role"]
        content = turn_content["content"]
        content = content.replace(
            "Available Tools (already imported for you):\n```\n", "```python\n\n"
        )
        content = content.replace("\n", "<br>")

        st.markdown(
            f"##### ROLE: <span style='color:{dict(ROLES)[role]}'>{role}</span>",
            unsafe_allow_html=True,
        )

        if turn_id == 0:
            # add newline before and after the code block
            # remove the <br> inside code block using re
            content = re.sub(
                r"```(.*?)```",
                lambda m: "\n" + m.group(0).replace("<br>", "\n") + "\n",
                content,
            )
        # use color to highlight span of regex by using re.sub
        for regex, color, field_name in REGEX_TO_HIGHLIGHT:
            content = re.sub(
                regex,
                # some hacky way to make sure the html-style prompt is escaped
                f"\n<span style='color:{color}'>&lt;{field_name}&gt; <b>\\1</b> &lt;/{field_name}&gt;</span>\n"
                if field_name != "execute"
                else lambda m: (
                    f"<span style='color:{color}'> &lt;{field_name}&gt; </span>\n"
                    + "```python\n{}\n```\n\n".format(m.group(1).replace("<br>", "\n"))
                    + f"<span style='color:{color}'> &lt;/{field_name}&gt; </span>\n"
                ),
                content,
            )
        st.write(content, unsafe_allow_html=True)
        # st.text(content)


visualize_row(row_dict)

scripts/visualizer.py
xingyaoww/mint-bench
from setuptools import setup, find_packages

with open("requirements.txt", "r") as f:
    requirements = f.read().splitlines()

setup(
    name="mint-bench",
    version="1.0.1",
    author="Xingyao Wang",
    author_email="xingyao6@illinois.edu",
    description="MINT Benchmark to Evaluate LLMs' Multi-Turn Interaction Capabilities.",
    url="https://xingyaoww.github.io/mint-bench",
    packages=find_packages(),
    install_requires=requirements,  # List of requirements read from requirements.txt
    classifiers=[
        "Programming Language :: Python :: 3",
    ],
    python_requires=">=3.8",
)

setup.py
xingyaoww/mint-bench