from typing import Any
from pandasai.pipelines.base_logic_unit import BaseLogicUnit
from pandasai.responses.context import Context
from pandasai.responses.response_parser import IResponseParser, ResponseParser


class ProcessOutput(BaseLogicUnit):
    """
    Executes the code generated by the prompt
    """

    _response_parser: IResponseParser

    def __init__(self, response_parser: IResponseParser = ResponseParser):
        self._response_parser = response_parser

    def execute(self, input: Any, **kwargs) -> Any:
        dfs = kwargs["context"].dfs

        context = Context(kwargs["config"], kwargs["logger"], dfs[0].engine)

        return self._response_parser(context).parse(input)
