import re
from typing import Union

from langchain.schema import AgentAction, AgentFinish, OutputParserException

from src.agents.agent import AgentOutputParser


class ReActOutputParser(AgentOutputParser):
    """Output parser for the ReAct agent."""

    def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
        action_prefix = "Action: "
        if not text.strip().split("\n")[-1].startswith(action_prefix):
            raise OutputParserException(f"Could not parse LLM Output: {text}")
        action_block = text.strip().split("\n")[-1]

        action_str = action_block[len(action_prefix) :]
        # Parse out the action and the directive.
        re_matches = re.search(r"(.*?)\[(.*?)\]", action_str)
        if re_matches is None:
            raise OutputParserException(
                f"Could not parse action directive: {action_str}"
            )
        action, action_input = re_matches.group(1), re_matches.group(2)
        if action == "Finish":
            return AgentFinish({"output": action_input}, text)
        else:
            return AgentAction(action, action_input, text)

    @property
    def _type(self) -> str:
        return "react"

langchain.schema.AgentFinish langchain.schema.AgentAction langchain.schema.OutputParserException
[((685, 726), 're.search', 're.search', (['"""(.*?)\\\\[(.*?)\\\\]"""', 'action_str'], {}), "('(.*?)\\\\[(.*?)\\\\]', action_str)\n", (694, 726), False, 'import re\n'), ((444, 504), 'langchain.schema.OutputParserException', 'OutputParserException', (['f"""Could not parse LLM Output: {text}"""'], {}), "(f'Could not parse LLM Output: {text}')\n", (465, 504), False, 'from langchain.schema import AgentAction, AgentFinish, OutputParserException\n'), ((775, 847), 'langchain.schema.OutputParserException', 'OutputParserException', (['f"""Could not parse action directive: {action_str}"""'], {}), "(f'Could not parse action directive: {action_str}')\n", (796, 847), False, 'from langchain.schema import AgentAction, AgentFinish, OutputParserException\n'), ((1000, 1043), 'langchain.schema.AgentFinish', 'AgentFinish', (["{'output': action_input}", 'text'], {}), "({'output': action_input}, text)\n", (1011, 1043), False, 'from langchain.schema import AgentAction, AgentFinish, OutputParserException\n'), ((1077, 1116), 'langchain.schema.AgentAction', 'AgentAction', (['action', 'action_input', 'text'], {}), '(action, action_input, text)\n', (1088, 1116), False, 'from langchain.schema import AgentAction, AgentFinish, OutputParserException\n')]