import re
from langchain.agents import AgentOutputParser
from langchain.schema import AgentAction, AgentFinish, OutputParserException
from typing import Union

from cat.mad_hatter.mad_hatter import MadHatter
from cat.log import log


class ChooseProcedureOutputParser(AgentOutputParser):

    def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:

        # Check if agent should finish
        if "Final Answer:" in llm_output:
            return AgentFinish(
                # Return values is generally always a dictionary with a single `output` key
                # It is not recommended to try anything else at the moment :)
                return_values={"output": llm_output.split("Final Answer:")[-1].strip()},
                log=llm_output,
            )
        
        # Parse out the action and action input
        regex = r"Action\s*\d*\s*:(.*?)\nAction\s*\d*\s*Input\s*\d*\s*:[\s]*(.*)"
        match = re.search(regex, llm_output, re.DOTALL)
        if not match:
            raise OutputParserException(f"Could not parse LLM output: `{llm_output}`")
        
        # Extract action
        action = match.group(1).strip()
        action_input = match.group(2)

        if action == "none_of_the_others":
            return AgentFinish(
                # Return values is generally always a dictionary with a single `output` key
                # It is not recommended to try anything else at the moment :)
                return_values={"output": None},
                log=llm_output,
            )

        mh = MadHatter()

        for Form in mh.forms:
            if Form.name == action:
                return AgentFinish(
                    return_values={
                        "output": None,
                        "form": action
                    },
                    log=llm_output,
                )

        # Return the action and action input
        return AgentAction(tool=action, tool_input=action_input.strip(" ").strip('"'), log=llm_output)
langchain.schema.AgentFinish langchain.schema.OutputParserException
[((936, 975), 're.search', 're.search', (['regex', 'llm_output', 're.DOTALL'], {}), '(regex, llm_output, re.DOTALL)\n', (945, 975), False, 'import re\n'), ((1551, 1562), 'cat.mad_hatter.mad_hatter.MadHatter', 'MadHatter', ([], {}), '()\n', (1560, 1562), False, 'from cat.mad_hatter.mad_hatter import MadHatter\n'), ((1016, 1084), 'langchain.schema.OutputParserException', 'OutputParserException', (['f"""Could not parse LLM output: `{llm_output}`"""'], {}), "(f'Could not parse LLM output: `{llm_output}`')\n", (1037, 1084), False, 'from langchain.schema import AgentAction, AgentFinish, OutputParserException\n'), ((1260, 1319), 'langchain.schema.AgentFinish', 'AgentFinish', ([], {'return_values': "{'output': None}", 'log': 'llm_output'}), "(return_values={'output': None}, log=llm_output)\n", (1271, 1319), False, 'from langchain.schema import AgentAction, AgentFinish, OutputParserException\n'), ((1653, 1728), 'langchain.schema.AgentFinish', 'AgentFinish', ([], {'return_values': "{'output': None, 'form': action}", 'log': 'llm_output'}), "(return_values={'output': None, 'form': action}, log=llm_output)\n", (1664, 1728), False, 'from langchain.schema import AgentAction, AgentFinish, OutputParserException\n')]