solver.py

"""
Codes related to CoTSolver
"""

from abc import ABC, abstractmethod
from typing import override
from agent import ChatBot

COT_AI_PROMPT = (
    "Let's first understand the problem, extract relevant variables and their corresponding numerals,"
    "and **devise a plan**. Then, let's **carry out the plan**, caculate intermediate results"
    "(pay attention to calculation and common sense), solve the problem step by step, and show the answer."
)


class CoTSolver(ABC):
    """
    Chain-of-Thought Solver
    """

    @abstractmethod
    def solve(self) -> str:
        """
        Solve the problem.

        Returns:
            - the final answer
        """

    @abstractmethod
    def set_problem(self, problem: str):
        """
        Set the problem description to be solved.
        NOTICE: The options should not be contained if it's a multiple choices problem.
        """

    @abstractmethod
    def send_human_message(self, msg: str) -> ChatBot:
        """
        get the agent
        """

    def solve_numerical(self) -> str:
        """
        Solve a numerical problem, e.g. return only a number.
        """
        self.solve()
        return self.send_human_message(
            "Therefore the answer is?"
            "Output only a real number(e.g., 3.14). Do not use fractional form(like 1/2 or 3/4)."
            "Do not show a equation like 1 + 1 = 2. In this case, output 2 only."
            "If your answer is a repeating decimal, round it to six decimal places."
        )

    def solve_multichoice(self, options: dict[str, str]) -> str:
        """
        Solve a multiple-choice question, e.g. return one of [A, B, C, D, E, ...]
        """
        self.solve()
        options_lines = []
        for k, v in options.items():
            options_lines.append(f"{k}: {v}")
        return self.send_human_message(
            f"Here are {len(options_lines)} options for the answer:\n"
            + "\n".join(options_lines)
            + "Please choose and output one of the upper letter of the options, e.g. A. "
            "If you think none of them is correct, just output the most likely letter."
            "Do not output anything beside the letter like A"
            "Please do not use bold formatting in your output."
        )

    def solve_boolean(self, boolean_expression: tuple[str, str]):
        """
        Args:
            boolean_expression: two str in a tuple, presenting true and false. e.g. ["True", "False"]
        """
        self.solve()
        return self.send_human_message(
            "Decide the answer to the problem to be true or false."
            f"If you think it's true, output {boolean_expression[0]},"
            f"else output {boolean_expression[1]}"
            f"Do not output anything beside the judgment like f{boolean_expression[0]}."
        )


class ZSCoTSolver(CoTSolver):
    """
    Wraps an agent and use zero-shot CoT to solve a problem(usually math).
    """

    def __init__(self, problem: str = None, model_name: str = None):
        self._problem = problem
        self._agent = ChatBot(model_name=model_name)

    @override
    def send_human_message(self, msg):
        return self._agent.post_human(msg)

    def set_problem(self, problem: str) -> None:
        """
        Simple Setter
        """
        self._problem = problem

    @override
    def solve(self) -> str:
        """
        Solve the problem.

        Returns:
            - a number indicates the final answer
        """
        self._agent.clear_history()
        self._agent.store_human(self._problem)
        return self._agent.post_ai("Let's think step by step.")


class PSCoTSolver(CoTSolver):
    """
    Wraps an agent and use Plan-and-Solve CoT to solve a problem(usually math).
    """

    def __init__(self, problem: str = None, model_name: str = None):
        self._problem = problem
        self._agent = ChatBot(model_name=model_name)

    @override
    def send_human_message(self, msg):
        return self._agent.post_human(msg)

    def set_problem(self, problem: str):
        """
        Simple Setter
        """
        self._problem = problem

    def solve(self) -> str:
        """
        Solve the problem.

        Returns:
            - a number indicates the final answer
        """
        self._agent.clear_history()
        self._agent.store_human(self._problem)
        return self._agent.post_ai(COT_AI_PROMPT)


class GiveAListSolver(CoTSolver):
    """
    Prompot the agent to give a list of steps to solve the problem
    """

    def __init__(self, problem: str = None, model_name: str = None):
        self._problem = problem
        self._agent = ChatBot(
            system_prompt=(
                "You are going to solve a problem. "
                "First, you should extract variables and their corresponding numerals."
                "Second, make a list. Each element of the list is a track of how these variables evolve."
                "The variable and its corresponding numerals should be included in the list."
                "After each manipulation in the problem, the variable may change(or not)."
                "You should clearly describe how the variable changes in the list."
                "Finally, the last element of the list should be the final answer of the problem."
            ),
            model_name=model_name,
        )

    @override
    def send_human_message(self, msg):
        return self._agent.post_human(msg)

    def set_problem(self, problem: str):
        """
        Simple Setter
        """
        self._problem = problem

    def solve(self) -> str:
        """
        Prompt the agent to give a list of steps to solve the problem
        """
        self._agent.clear_history()
        self._agent.store_human(self._problem)
        return self._agent.post_ai(
            "Ok, I will think step by step and give you a list of steps to solve the problem.\n"
            "1. "
        )