from typing import Union
import re

from langchain.agents import AgentOutputParser
from langchain.schema import AgentAction, AgentFinish


class CustomOutputParser(AgentOutputParser):
    """
    This is the output parser for the multi-tool agent. It parses the output from the LLM model.
    """

    def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:
        """
        This function is used to parse the output from the LLM model. It checks if the output is the final answer or an action.
        """

        if "Final Answer:" in llm_output:
            return AgentFinish(
                return_values={"output": llm_output.split("Final Answer:")[-1].strip()},
                log=llm_output,
            )

        regex = r"Action: (.*?)[\n]*Action Input:[\s]*(.*)"
        match = re.search(regex, llm_output, re.DOTALL)

        if not match:
            raise ValueError(f"Could not parse LLM output: `{llm_output}`")
        action = match.group(1).strip()
        action_input = match.group(2)

        return AgentAction(
            tool=action, tool_input=action_input.strip(" ").strip('"'), log=llm_output
        )