                            msg_head="download file success").debug(path)
                return send_file(path, as_attachment=True)

        logger.bind(user_id=user_id, api="/download",
                    msg_head="download file failed").debug(path)
        return Response(response=None,
                        status=f"{INTERNAL} Download file failed: file not correctlt sent")

    except Exception as e:
        print(str(e))
        import traceback
        traceback.print_exc()

        logger.bind(user_id=user_id, api="/download",
                    msg_head="download file failed").error(str(e))

        return Response(response=None,
                        status=f"{INTERNAL} Download file failed: {str(e)}")

def _generate_directory_name(name: str, x:int=0) -> Any:
    dir_name = (name + ("_" + str(x) if x != 0 else "")).strip()
    if not os.path.exists(dir_name):
        return dir_name
    else:
        return _generate_directory_name(name, x + 1)

@app.route("/api/file_system/create_folder", methods=["POST"])
def create_folders() -> Response:
    """Creates a folder in the filesystem."""
    request_json = request.get_json()
    (user_id, chat_id) = get_user_and_chat_id_from_request_json(request_json)
    root_path = create_personal_folder(user_id)
    if os.path.exists(root_path) and os.path.isdir(root_path):
        try:
            new_path = _generate_directory_name(os.path.join(root_path, "Folder"))
            os.makedirs(new_path, exist_ok=False)

            logger.bind(
                user_id=user_id, chat_id=chat_id, api="/create_folder",
                msg_head="Create folder success"
            ).debug(new_path)

            return jsonify({"success": True, "message": "Folder created successfully"})
        except Exception as e:


        print(str(e))
        import traceback
        traceback.print_exc()

        logger.bind(user_id=user_id, api="/download",
                    msg_head="download file failed").error(str(e))

        return Response(response=None,
                        status=f"{INTERNAL} Download file failed: {str(e)}")

def _generate_directory_name(name: str, x:int=0) -> Any:
    dir_name = (name + ("_" + str(x) if x != 0 else "")).strip()
    if not os.path.exists(dir_name):
        return dir_name
    else:
        return _generate_directory_name(name, x + 1)

@app.route("/api/file_system/create_folder", methods=["POST"])
def create_folders() -> Response:
    """Creates a folder in the filesystem."""
    request_json = request.get_json()
    (user_id, chat_id) = get_user_and_chat_id_from_request_json(request_json)
    root_path = create_personal_folder(user_id)
    if os.path.exists(root_path) and os.path.isdir(root_path):
        try:
            new_path = _generate_directory_name(os.path.join(root_path, "Folder"))
            os.makedirs(new_path, exist_ok=False)

            logger.bind(
                user_id=user_id, chat_id=chat_id, api="/create_folder",
                msg_head="Create folder success"
            ).debug(new_path)

            return jsonify({"success": True, "message": "Folder created successfully"})
        except Exception as e:
            logger.bind(user_id=user_id, chat_id=chat_id, api="/create_folder",
                        msg_head="Create folder failed").error(
                str(e)
            )

            return jsonify({"success": False, "message": str(e)})
    else:
        logger.bind(user_id=user_id, chat_id=chat_id, api="/create_folder",
                    msg_head="Create folder failed").error(
            "Root path does not exist."
        )

        return Response(response=None, status=f"{INTERNAL} Root path does not exist")

@app.route("/api/file_system/update", methods=["POST"])
def rename_folder() -> Response:
    """Renames a folder in the filesystem."""
    request_json = request.get_json()
    (user_id, chat_id) = get_user_and_chat_id_from_request_json(request_json)
    root_path = create_personal_folder(user_id)
    node = request_json["node"]
    rename_value = request_json["rename_value"]
    if os.path.exists(root_path) and os.path.isdir(root_path):
        try:
            current_path_tree_list: list = []
            id_to_path_dict = {0: root_path}
            _path_tree_for_react_dnd_treeview(current_path_tree_list, id_to_path_dict,
                                              root_path, 0)
            path = id_to_path_dict[node["id"]]
            new_path = os.path.join(os.path.dirname(path), rename_value)
            shutil.move(path, new_path)


            ).debug(new_path)

            return jsonify({"success": True, "message": "Folder created successfully"})
        except Exception as e:
            logger.bind(user_id=user_id, chat_id=chat_id, api="/create_folder",
                        msg_head="Create folder failed").error(
                str(e)
            )

            return jsonify({"success": False, "message": str(e)})
    else:
        logger.bind(user_id=user_id, chat_id=chat_id, api="/create_folder",
                    msg_head="Create folder failed").error(
            "Root path does not exist."
        )

        return Response(response=None, status=f"{INTERNAL} Root path does not exist")

@app.route("/api/file_system/update", methods=["POST"])
def rename_folder() -> Response:
    """Renames a folder in the filesystem."""
    request_json = request.get_json()
    (user_id, chat_id) = get_user_and_chat_id_from_request_json(request_json)
    root_path = create_personal_folder(user_id)
    node = request_json["node"]
    rename_value = request_json["rename_value"]
    if os.path.exists(root_path) and os.path.isdir(root_path):
        try:
            current_path_tree_list: list = []
            id_to_path_dict = {0: root_path}
            _path_tree_for_react_dnd_treeview(current_path_tree_list, id_to_path_dict,
                                              root_path, 0)
            path = id_to_path_dict[node["id"]]
            new_path = os.path.join(os.path.dirname(path), rename_value)
            shutil.move(path, new_path)

            logger.bind(user_id=user_id, chat_id=chat_id, api="/update",
                        msg_head="Rename folder success").debug(
                f"{path} to {new_path}"
            )

            return jsonify({"success": True, "message": "Folder created successfully"})
        except Exception as e:
            logger.bind(user_id=user_id, chat_id=chat_id, api="/update",
                        msg_head="Rename folder failed").error(str(e))

            return jsonify({"success": False, "message": str(e)})
    else:
        logger.bind(user_id=user_id, chat_id=chat_id, api="/update",
                    msg_head="Rename folder failed").error(
            "Root path does not exist."
        )

        return Response(response=None, status=f"{INTERNAL} Root path does not exist")

@app.route("/api/file_system/get_path_tree", methods=["POST"])
def get_path_tree() -> Response:
    """Gets a file path tree of one file."""
    try:
        request_json = request.get_json()
        user_id = request_json.pop("user_id", DEFAULT_USER_ID)
        if user_id == "":  # front-end may enter empty user_id
            return jsonify([])
        root_path = create_personal_folder(user_id)
        highlighted_files = request_json.get("highlighted_files", [])
        if root_path is None:
            return {"error": "root_path parameter is required", "error_code": 404}
        if os.path.exists(root_path) and os.path.isdir(root_path):
            current_path_tree_list: list = []
            id_to_path_dict = {0: root_path}
            _path_tree_for_react_dnd_treeview(current_path_tree_list, id_to_path_dict,
                                              root_path, 0,
                                              highlighted_files=highlighted_files)


                        msg_head="Rename folder success").debug(
                f"{path} to {new_path}"
            )

            return jsonify({"success": True, "message": "Folder created successfully"})
        except Exception as e:
            logger.bind(user_id=user_id, chat_id=chat_id, api="/update",
                        msg_head="Rename folder failed").error(str(e))

            return jsonify({"success": False, "message": str(e)})
    else:
        logger.bind(user_id=user_id, chat_id=chat_id, api="/update",
                    msg_head="Rename folder failed").error(
            "Root path does not exist."
        )

        return Response(response=None, status=f"{INTERNAL} Root path does not exist")

@app.route("/api/file_system/get_path_tree", methods=["POST"])
def get_path_tree() -> Response:
    """Gets a file path tree of one file."""
    try:
        request_json = request.get_json()
        user_id = request_json.pop("user_id", DEFAULT_USER_ID)
        if user_id == "":  # front-end may enter empty user_id
            return jsonify([])
        root_path = create_personal_folder(user_id)
        highlighted_files = request_json.get("highlighted_files", [])
        if root_path is None:
            return {"error": "root_path parameter is required", "error_code": 404}
        if os.path.exists(root_path) and os.path.isdir(root_path):
            current_path_tree_list: list = []
            id_to_path_dict = {0: root_path}
            _path_tree_for_react_dnd_treeview(current_path_tree_list, id_to_path_dict,
                                              root_path, 0,
                                              highlighted_files=highlighted_files)
            return jsonify(current_path_tree_list)
        else:
            return Response(response=None, status=f"{UNFOUND} Directory not found")
    except Exception as e:
        return Response(response=None, status=f"{INTERNAL} Directory not found")

@app.route("/api/set_default_examples", methods=["POST"])
def set_default_examples() -> Response:
    """Sets default files for each user."""
    try:
        # Should be called after auth is verified
        request_json = request.get_json()
        user_id = request_json.pop("user_id", DEFAULT_USER_ID)
        root_path = create_personal_folder(user_id)
        example_dir = os.path.dirname(os.path.dirname(app.config["UPLOAD_FOLDER"]))
        example_path = os.path.join(example_dir, "data/examples/")
        if os.path.exists(example_path):
            shutil.copytree(example_path, root_path, dirs_exist_ok=True)
            return jsonify(
                {"success": True, "message": "Default examples are set successfully"})
        else:
            return Response(response=None,
                            status=f"{UNFOUND} Directory not found at {example_dir}")
    except Exception as e:


        if user_id == "":  # front-end may enter empty user_id
            return jsonify([])
        root_path = create_personal_folder(user_id)
        highlighted_files = request_json.get("highlighted_files", [])
        if root_path is None:
            return {"error": "root_path parameter is required", "error_code": 404}
        if os.path.exists(root_path) and os.path.isdir(root_path):
            current_path_tree_list: list = []
            id_to_path_dict = {0: root_path}
            _path_tree_for_react_dnd_treeview(current_path_tree_list, id_to_path_dict,
                                              root_path, 0,
                                              highlighted_files=highlighted_files)
            return jsonify(current_path_tree_list)
        else:
            return Response(response=None, status=f"{UNFOUND} Directory not found")
    except Exception as e:
        return Response(response=None, status=f"{INTERNAL} Directory not found")

@app.route("/api/set_default_examples", methods=["POST"])
def set_default_examples() -> Response:
    """Sets default files for each user."""
    try:
        # Should be called after auth is verified
        request_json = request.get_json()
        user_id = request_json.pop("user_id", DEFAULT_USER_ID)
        root_path = create_personal_folder(user_id)
        example_dir = os.path.dirname(os.path.dirname(app.config["UPLOAD_FOLDER"]))
        example_path = os.path.join(example_dir, "data/examples/")
        if os.path.exists(example_path):
            shutil.copytree(example_path, root_path, dirs_exist_ok=True)
            return jsonify(
                {"success": True, "message": "Default examples are set successfully"})
        else:
            return Response(response=None,
                            status=f"{UNFOUND} Directory not found at {example_dir}")
    except Exception as e:
        return Response(response=None,
                        status=f"{INTERNAL} Fail to Set Default Examples")


# flake8: noqa

PREFIX = """You are XLang Agent , a friendly and intuitive interface developed by the XLang Team to guide human through every stage of human data lifecycle. Whether human are loading, processing, or interpreting data, XLang Agent is always at human's fingertips through our interactive chat system.

Empowered by an array of innovative tools that can generate and execute code, XLang Agent delivers robust, reliable answers to human queries. Whenever possible, You employs these tools to give human rich insights, like dynamic code generation & execution and compelling visualizations. And You will always proactively and correctly using all tools to help with human.

Get ready for a seamless and insightful journey with XLang Agent, the personal assistant for all things data!

TOOLS
------
You have direct access to following tools.
"""

FORMAT_INSTRUCTIONS = """RESPONSE FORMAT INSTRUCTIONS
----------------------------

When you use tools or generate final answer, please output a response in one of two formats:
**Option 1: Explain and Use Tool**
If the response involves using a tool, you can start with a natural language explanation[Optional], plus exactly one tool calling[MUST]. But **make sure no any words & answer appended after tool calling json**. The tool calling format should be a markdown code snippet with the following JSON schema:

```json
{{{{
    "action": string wrapped with \"\", // The action to take. Must be one in the list [{tool_names}]
    "action_input": string wrapped with \"\" // Natural language query to be input to the action tool.
}}}}
```

[**Restriction**] Please note that ONLY one tool should be used per round, and you MUST stop generating right after tool calling and make sure no any text appended after tool calling markdown code snippet. Save your words.

NEVER EVER EVER make up a tool not in [{tool_names}]
NEVER EVER EVER generate code as action input when using tool. Just input natural language by using/paraphrasing human query.

**Option #2:**
Use this if you want to respond directly to the human.
If you want to respond directly to the human without using a tool, provide a plain natural language response. However, if you initially generated a natural language response and then decide to use a tool, make sure to include the tool action and input after the initial response.

Note if the human asks for malicious code, and just respond directly to deny the request and give your professional reason. Don't use any tool.
The malicious code includes but not limited to:
1. Endless operations and excessive waiting  (e.g., while True, long print, input())
2. System crash (e.g., any risky system command)
3. Data loss (e.g., list or delete files)
4. Leak sensitive information (e.g., os.getenv())
5. Establish network connections (e.g., requests.get())
6. Cause any other security issues

[Mandatory to notice] It is imperative and a must to utilize tools whenever the human's query tasks that implies using tools, such as searching online, generating code, executing code, or any other complex functionalities. You must try to use tools to solve human queries in these cases.

Begin.
"""

SUFFIX = """{input}"""

TEMPLATE_TOOL_RESPONSE = """TOOL RESPONSE:
---------------------
{observation}

THOUGHT
--------------------

Okay, So what's next? Let's assess if the tool response is enough to answer the human's initial query. Please follow these instructions:

1. Evaluate Tool Response [Mandatory]: Carefully evaluate the tool's response and determine if it sufficiently addresses the human's query. Consider the content and implications of the tool's response.

2. Consider Additional Tool Use [Optional 2 or 3]: If the tool response does not fully address the query or if an error occurred during execution, you may proceed with additional tool usage. However, exercise caution and limit the number of iterations to a maximum of three. You can start with a natural language explanation[Optional], plus exactly one tool calling[MUST]. But **make sure no any words & answer appended after tool calling json**. Follow this format for additional tool usage:

```json
{{{{
    "action": string wrapped with \"\", // The action to take. Must be one of [{tool_names}]
    "action_input": string wrapped with \"\" // Natural language query to be input to the action tool
}}}}
```
[**Restriction**] Please note that only one tool should be used per round, and you MUST stop generating right after tool calling and make sure no any text appended after tool calling markdown code snippet.

3. Deliver Comprehensive Answer [Optional 2 or 3]: If the tool response sufficiently addresses the query, deliver a comprehensive answer to the human. Focus solely on the content and implications of the tool's response. MUST NOT include explanations of the tool's functions.

3.1. Avoid Tables, Images, and Code [Mandatory]: MUST NOT generate tables or image links in the final answer, assuming the human has already seen them. Avoid generating code in the final answer as well. Instead, paraphrase the code into a human query if you need to explain it.

Note. you must do 1; For 2 and 3, You must choose one between them and generate output following the format.

Begin.
"""

# models like anthropic claude-v1 or claude-2 can only return valid completion with human message as the last message, so we append the fake AI message at the end.
fake_continue_prompt = {
    "claude-2": "you can start to think and respond to me using the above formats. No Apology. Just respond with format in Option 2(use tool) or Option 3(direct text response), no other words.\n\nBegin.",
    "claude-v1": "you can start to think and respond to me using the above formats. No Apology. Just respond with format in Option 2(use tool) or Option 3(direct text response), no other words.\n\nBegin.",
}


"""An agent designed to hold a conversation in addition to using tools."""
from __future__ import annotations

from typing import Any, List, Optional, Sequence, Tuple, Union
from typing_extensions import override
from pydantic import Field

from langchain.agents.agent import AgentOutputParser
from langchain.agents.utils import validate_tools_single_input
from langchain.base_language import BaseLanguageModel
from langchain.callbacks.manager import AsyncCallbackManagerForToolRun, CallbackManagerForToolRun, Callbacks
from langchain.chains import LLMChain
from langchain.schema import AgentAction, AgentFinish, HumanMessage, AIMessage, BaseMessage, BaseOutputParser
from langchain.tools.base import BaseTool

from real_agents.adapters.agent_helpers.agent import Agent
from real_agents.adapters.agent_helpers.output_parser import ConversationOutputParser
from real_agents.data_agent.copilot_prompt import PREFIX, SUFFIX, TEMPLATE_TOOL_RESPONSE, fake_continue_prompt
from real_agents.adapters.data_model import DataModel, MessageDataModel
from langchain.prompts import (
    BasePromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
)

class ExceptionTool(BaseTool):
    name = "_Exception"
    description = "Exception tool"

    def _run(
        self,
        query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        return query

    async def _arun(
        self,
        query: str,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> str:
        return query

class ConversationalChatAgent(Agent):
    """An agent designed to hold a conversation in addition to using data tools."""

    output_parser: ConversationOutputParser = Field(default_factory=ConversationOutputParser())
    template_tool_response: str = TEMPLATE_TOOL_RESPONSE
    continue_model: Optional[str] = None

    @classmethod
    def _get_default_output_parser(cls, **kwargs: Any) -> ConversationOutputParser:
        return ConversationOutputParser()

    @property
    def _agent_type(self) -> str:
        raise NotImplementedError

    @property
    def observation_prefix(self) -> str:
        """Prefix to append the observation with."""
        return "Observation: "

    @property
    def llm_prefix(self) -> str:
        """Prefix to append the llm call with."""
        return "Thought:"

    @classmethod
    def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
        super()._validate_tools(tools)
        validate_tools_single_input(cls.__name__, tools)

    @classmethod
    def create_prompt(
        cls,
        tools: Sequence[BaseTool],
        system_message: str = PREFIX,
        human_message: str = SUFFIX,
        input_variables: Optional[List[str]] = None,
        output_parser: Optional[BaseOutputParser] = None,
    ) -> BasePromptTemplate:
        # tools
        tool_strings = "\n".join([f"> {tool.name}: {tool.description}" for tool in tools])
        tool_names = ", ".join([tool.name for tool in tools])
        _output_parser = output_parser or cls._get_default_output_parser()

        # format instructions for system message
        format_instructions = _output_parser.get_format_instructions()
        format_instructions = format_instructions.format(tool_names=tool_names)

        # system message
        system_message = system_message + f"{tool_strings}\n\n{format_instructions}"

        # human input
        final_prompt = human_message
        if input_variables is None:
            input_variables = ["input", "chat_history", "agent_scratchpad"]
        messages = [
            SystemMessagePromptTemplate.from_template(system_message),
            MessagesPlaceholder(variable_name="chat_history"),
            HumanMessagePromptTemplate.from_template(final_prompt),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]
        return ChatPromptTemplate(input_variables=input_variables, messages=messages)

    @override
    def _construct_scratchpad(self, intermediate_steps: List[Tuple[AgentAction, str]]) -> List[BaseMessage]:
        """Construct the scratchpad that lets the agent continue its thought process."""
        thoughts: List[BaseMessage] = []

        # Try to only use AI message for scratchpad
        content = []
        for idx, (action, full_observation) in enumerate(intermediate_steps):
            content.append(MessageDataModel.extract_action_for_llm(action.log))

            observation = full_observation
            if isinstance(full_observation, DataModel):
                llm_raw_observation = full_observation.get_llm_side_data()
                observation = MessageDataModel.extract_tool_response_for_llm(llm_raw_observation)
                tool_response = self.template_tool_response.format(
                    observation=str(observation), tool_names=self.allowed_tools
                )
                if idx == len(intermediate_steps) - 1:
                    content.append(tool_response)
                else:
                    content.append(observation)
        content_str = "\n".join(content)
        thoughts.append(AIMessage(content=content_str))
        if self.continue_model is not None and len(intermediate_steps) != 0:
            thoughts.append(HumanMessage(content=fake_continue_prompt[self.continue_model]))
        return thoughts

    @override
    def plan(
        self,
        intermediate_steps: List[Tuple[AgentAction, str]],
        callbacks: Callbacks = None,
        **kwargs: Any,
    ) -> Union[AgentAction, AgentFinish]:
        """Given input, decided what to do.

        Args:
            intermediate_steps: Steps the LLM has taken to date,
                along with observations
            callbacks: Callbacks to run.
            **kwargs: User inputs.

        Returns:
            Action specifying what tool to use.
        """
        full_inputs = self.get_full_inputs(intermediate_steps, **kwargs)
        system_prompt = self.llm_chain.prompt.messages[0].format().content
        system_prompt_tokens = MessageDataModel._count_tokens(system_prompt)
        max_tokens = 8000
        max_gen_tokens = 1000
        # FIXME: need more accurate token limit calculation
        full_inputs = MessageDataModel.truncate_chat_history(
            full_inputs, max_token=max_tokens - system_prompt_tokens - max_gen_tokens
        )
        full_output = self.llm_chain.predict(callbacks=callbacks, **full_inputs)

        return self.output_parser.parse(full_output)

    @classmethod
    def from_llm_and_tools(
        cls,
        llm: BaseLanguageModel,
        tools: Sequence[BaseTool],
        callbacks: Callbacks = None,
        output_parser: Optional[AgentOutputParser] = None,
        system_message: str = PREFIX,
        human_message: str = SUFFIX,
        input_variables: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Agent:
        """Construct an agent from an LLM and tools."""
        cls._validate_tools(tools)

        _output_parser = output_parser or cls._get_default_output_parser()
        prompt = cls.create_prompt(
            tools,
            system_message=system_message,
            human_message=human_message,
            input_variables=input_variables,
            output_parser=_output_parser,
        )
        llm_chain = LLMChain(
            llm=llm,
            prompt=prompt,
        )
        tool_names = [tool.name for tool in tools]
        return cls(
            llm_chain=llm_chain,
            allowed_tools=tool_names,
            output_parser=_output_parser,
            **kwargs,
        )


ExceptionTool

from langchain.agents.utils import validate_tools_single_input
from langchain.base_language import BaseLanguageModel
from langchain.callbacks.manager import AsyncCallbackManagerForToolRun, CallbackManagerForToolRun, Callbacks
from langchain.chains import LLMChain
from langchain.schema import AgentAction, AgentFinish, HumanMessage, AIMessage, BaseMessage, BaseOutputParser
from langchain.tools.base import BaseTool

from real_agents.adapters.agent_helpers.agent import Agent
from real_agents.adapters.agent_helpers.output_parser import ConversationOutputParser
from real_agents.data_agent.copilot_prompt import PREFIX, SUFFIX, TEMPLATE_TOOL_RESPONSE, fake_continue_prompt
from real_agents.adapters.data_model import DataModel, MessageDataModel
from langchain.prompts import (
    BasePromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
)

class ExceptionTool(BaseTool):
    name = "_Exception"
    description = "Exception tool"

    def _run(
        self,
        query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        return query

    async def _arun(
        self,
        query: str,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> str:
        return query

class ConversationalChatAgent(Agent):
    """An agent designed to hold a conversation in addition to using data tools."""

    output_parser: ConversationOutputParser = Field(default_factory=ConversationOutputParser())
    template_tool_response: str = TEMPLATE_TOOL_RESPONSE
    continue_model: Optional[str] = None

    @classmethod
    def _get_default_output_parser(cls, **kwargs: Any) -> ConversationOutputParser:
        return ConversationOutputParser()

    @property
    def _agent_type(self) -> str:
        raise NotImplementedError

    @property
    def observation_prefix(self) -> str:
        """Prefix to append the observation with."""


ConversationalChatAgent

class ExceptionTool(BaseTool):
    name = "_Exception"
    description = "Exception tool"

    def _run(
        self,
        query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        return query

    async def _arun(
        self,
        query: str,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> str:
        return query

class ConversationalChatAgent(Agent):
    """An agent designed to hold a conversation in addition to using data tools."""

    output_parser: ConversationOutputParser = Field(default_factory=ConversationOutputParser())
    template_tool_response: str = TEMPLATE_TOOL_RESPONSE
    continue_model: Optional[str] = None

    @classmethod
    def _get_default_output_parser(cls, **kwargs: Any) -> ConversationOutputParser:
        return ConversationOutputParser()

    @property
    def _agent_type(self) -> str:
        raise NotImplementedError

    @property
    def observation_prefix(self) -> str:
        """Prefix to append the observation with."""
        return "Observation: "

    @property
    def llm_prefix(self) -> str:
        """Prefix to append the llm call with."""
        return "Thought:"

    @classmethod
    def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
        super()._validate_tools(tools)
        validate_tools_single_input(cls.__name__, tools)

    @classmethod
    def create_prompt(
        cls,
        tools: Sequence[BaseTool],
        system_message: str = PREFIX,
        human_message: str = SUFFIX,
        input_variables: Optional[List[str]] = None,
        output_parser: Optional[BaseOutputParser] = None,
    ) -> BasePromptTemplate:
        # tools
        tool_strings = "\n".join([f"> {tool.name}: {tool.description}" for tool in tools])
        tool_names = ", ".join([tool.name for tool in tools])
        _output_parser = output_parser or cls._get_default_output_parser()

        # format instructions for system message
        format_instructions = _output_parser.get_format_instructions()
        format_instructions = format_instructions.format(tool_names=tool_names)

        # system message
        system_message = system_message + f"{tool_strings}\n\n{format_instructions}"

        # human input
        final_prompt = human_message
        if input_variables is None:
            input_variables = ["input", "chat_history", "agent_scratchpad"]
        messages = [
            SystemMessagePromptTemplate.from_template(system_message),
            MessagesPlaceholder(variable_name="chat_history"),
            HumanMessagePromptTemplate.from_template(final_prompt),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]
        return ChatPromptTemplate(input_variables=input_variables, messages=messages)

    @override
    def _construct_scratchpad(self, intermediate_steps: List[Tuple[AgentAction, str]]) -> List[BaseMessage]:
        """Construct the scratchpad that lets the agent continue its thought process."""
        thoughts: List[BaseMessage] = []

        # Try to only use AI message for scratchpad
        content = []
        for idx, (action, full_observation) in enumerate(intermediate_steps):
            content.append(MessageDataModel.extract_action_for_llm(action.log))

            observation = full_observation
            if isinstance(full_observation, DataModel):
                llm_raw_observation = full_observation.get_llm_side_data()
                observation = MessageDataModel.extract_tool_response_for_llm(llm_raw_observation)
                tool_response = self.template_tool_response.format(
                    observation=str(observation), tool_names=self.allowed_tools
                )
                if idx == len(intermediate_steps) - 1:
                    content.append(tool_response)
                else:
                    content.append(observation)
        content_str = "\n".join(content)
        thoughts.append(AIMessage(content=content_str))
        if self.continue_model is not None and len(intermediate_steps) != 0:
            thoughts.append(HumanMessage(content=fake_continue_prompt[self.continue_model]))
        return thoughts

    @override
    def plan(
        self,
        intermediate_steps: List[Tuple[AgentAction, str]],
        callbacks: Callbacks = None,
        **kwargs: Any,
    ) -> Union[AgentAction, AgentFinish]:
        """Given input, decided what to do.

        Args:
            intermediate_steps: Steps the LLM has taken to date,
                along with observations
            callbacks: Callbacks to run.
            **kwargs: User inputs.

        Returns:
            Action specifying what tool to use.
        """
        full_inputs = self.get_full_inputs(intermediate_steps, **kwargs)
        system_prompt = self.llm_chain.prompt.messages[0].format().content
        system_prompt_tokens = MessageDataModel._count_tokens(system_prompt)
        max_tokens = 8000
        max_gen_tokens = 1000
        # FIXME: need more accurate token limit calculation
        full_inputs = MessageDataModel.truncate_chat_history(
            full_inputs, max_token=max_tokens - system_prompt_tokens - max_gen_tokens
        )
        full_output = self.llm_chain.predict(callbacks=callbacks, **full_inputs)

        return self.output_parser.parse(full_output)

    @classmethod
    def from_llm_and_tools(
        cls,
        llm: BaseLanguageModel,
        tools: Sequence[BaseTool],
        callbacks: Callbacks = None,
        output_parser: Optional[AgentOutputParser] = None,
        system_message: str = PREFIX,
        human_message: str = SUFFIX,
        input_variables: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Agent:
        """Construct an agent from an LLM and tools."""
        cls._validate_tools(tools)

        _output_parser = output_parser or cls._get_default_output_parser()
        prompt = cls.create_prompt(
            tools,
            system_message=system_message,
            human_message=human_message,
            input_variables=input_variables,
            output_parser=_output_parser,
        )
        llm_chain = LLMChain(
            llm=llm,
            prompt=prompt,
        )
        tool_names = [tool.name for tool in tools]
        return cls(
            llm_chain=llm_chain,
            allowed_tools=tool_names,
            output_parser=_output_parser,
            **kwargs,
        )


from langchain.schema import AgentAction, AgentFinish, HumanMessage, AIMessage, BaseMessage, BaseOutputParser
from langchain.tools.base import BaseTool

from real_agents.adapters.agent_helpers.agent import Agent
from real_agents.adapters.agent_helpers.output_parser import ConversationOutputParser
from real_agents.data_agent.copilot_prompt import PREFIX, SUFFIX, TEMPLATE_TOOL_RESPONSE, fake_continue_prompt
from real_agents.adapters.data_model import DataModel, MessageDataModel
from langchain.prompts import (
    BasePromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
)

class ExceptionTool(BaseTool):
    name = "_Exception"
    description = "Exception tool"

    def _run(
        self,
        query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        return query

    async def _arun(
        self,
        query: str,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> str:
        return query

class ConversationalChatAgent(Agent):
    """An agent designed to hold a conversation in addition to using data tools."""

    output_parser: ConversationOutputParser = Field(default_factory=ConversationOutputParser())
    template_tool_response: str = TEMPLATE_TOOL_RESPONSE
    continue_model: Optional[str] = None

    @classmethod
    def _get_default_output_parser(cls, **kwargs: Any) -> ConversationOutputParser:
        return ConversationOutputParser()


from langchain.prompts import (
    BasePromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
)

class ExceptionTool(BaseTool):
    name = "_Exception"
    description = "Exception tool"

    def _run(
        self,
        query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        return query

    async def _arun(
        self,
        query: str,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> str:
        return query

class ConversationalChatAgent(Agent):
    """An agent designed to hold a conversation in addition to using data tools."""

    output_parser: ConversationOutputParser = Field(default_factory=ConversationOutputParser())
    template_tool_response: str = TEMPLATE_TOOL_RESPONSE
    continue_model: Optional[str] = None

    @classmethod
    def _get_default_output_parser(cls, **kwargs: Any) -> ConversationOutputParser:
        return ConversationOutputParser()

    @property
    def _agent_type(self) -> str:
        raise NotImplementedError

    @property
    def observation_prefix(self) -> str:
        """Prefix to append the observation with."""


        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        return query

    async def _arun(
        self,
        query: str,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> str:
        return query

class ConversationalChatAgent(Agent):
    """An agent designed to hold a conversation in addition to using data tools."""

    output_parser: ConversationOutputParser = Field(default_factory=ConversationOutputParser())
    template_tool_response: str = TEMPLATE_TOOL_RESPONSE
    continue_model: Optional[str] = None

    @classmethod
    def _get_default_output_parser(cls, **kwargs: Any) -> ConversationOutputParser:
        return ConversationOutputParser()

    @property
    def _agent_type(self) -> str:
        raise NotImplementedError

    @property
    def observation_prefix(self) -> str:
        """Prefix to append the observation with."""
        return "Observation: "

    @property
    def llm_prefix(self) -> str:
        """Prefix to append the llm call with."""
        return "Thought:"

    @classmethod
    def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
        super()._validate_tools(tools)
        validate_tools_single_input(cls.__name__, tools)


    async def _arun(
        self,
        query: str,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> str:
        return query

class ConversationalChatAgent(Agent):
    """An agent designed to hold a conversation in addition to using data tools."""

    output_parser: ConversationOutputParser = Field(default_factory=ConversationOutputParser())
    template_tool_response: str = TEMPLATE_TOOL_RESPONSE
    continue_model: Optional[str] = None

    @classmethod
    def _get_default_output_parser(cls, **kwargs: Any) -> ConversationOutputParser:
        return ConversationOutputParser()

    @property
    def _agent_type(self) -> str:
        raise NotImplementedError

    @property
    def observation_prefix(self) -> str:
        """Prefix to append the observation with."""
        return "Observation: "

    @property
    def llm_prefix(self) -> str:
        """Prefix to append the llm call with."""
        return "Thought:"

    @classmethod
    def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
        super()._validate_tools(tools)
        validate_tools_single_input(cls.__name__, tools)

    @classmethod
    def create_prompt(
        cls,
        tools: Sequence[BaseTool],


    ) -> str:
        return query

class ConversationalChatAgent(Agent):
    """An agent designed to hold a conversation in addition to using data tools."""

    output_parser: ConversationOutputParser = Field(default_factory=ConversationOutputParser())
    template_tool_response: str = TEMPLATE_TOOL_RESPONSE
    continue_model: Optional[str] = None

    @classmethod
    def _get_default_output_parser(cls, **kwargs: Any) -> ConversationOutputParser:
        return ConversationOutputParser()

    @property
    def _agent_type(self) -> str:
        raise NotImplementedError

    @property
    def observation_prefix(self) -> str:
        """Prefix to append the observation with."""
        return "Observation: "

    @property
    def llm_prefix(self) -> str:
        """Prefix to append the llm call with."""
        return "Thought:"

    @classmethod
    def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
        super()._validate_tools(tools)
        validate_tools_single_input(cls.__name__, tools)

    @classmethod
    def create_prompt(
        cls,
        tools: Sequence[BaseTool],
        system_message: str = PREFIX,
        human_message: str = SUFFIX,
        input_variables: Optional[List[str]] = None,
        output_parser: Optional[BaseOutputParser] = None,
    ) -> BasePromptTemplate:


    """An agent designed to hold a conversation in addition to using data tools."""

    output_parser: ConversationOutputParser = Field(default_factory=ConversationOutputParser())
    template_tool_response: str = TEMPLATE_TOOL_RESPONSE
    continue_model: Optional[str] = None

    @classmethod
    def _get_default_output_parser(cls, **kwargs: Any) -> ConversationOutputParser:
        return ConversationOutputParser()

    @property
    def _agent_type(self) -> str:
        raise NotImplementedError

    @property
    def observation_prefix(self) -> str:
        """Prefix to append the observation with."""
        return "Observation: "

    @property
    def llm_prefix(self) -> str:
        """Prefix to append the llm call with."""
        return "Thought:"

    @classmethod
    def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
        super()._validate_tools(tools)
        validate_tools_single_input(cls.__name__, tools)

    @classmethod
    def create_prompt(
        cls,
        tools: Sequence[BaseTool],
        system_message: str = PREFIX,
        human_message: str = SUFFIX,
        input_variables: Optional[List[str]] = None,
        output_parser: Optional[BaseOutputParser] = None,
    ) -> BasePromptTemplate:
        # tools
        tool_strings = "\n".join([f"> {tool.name}: {tool.description}" for tool in tools])
        tool_names = ", ".join([tool.name for tool in tools])
        _output_parser = output_parser or cls._get_default_output_parser()


    @classmethod
    def _get_default_output_parser(cls, **kwargs: Any) -> ConversationOutputParser:
        return ConversationOutputParser()

    @property
    def _agent_type(self) -> str:
        raise NotImplementedError

    @property
    def observation_prefix(self) -> str:
        """Prefix to append the observation with."""
        return "Observation: "

    @property
    def llm_prefix(self) -> str:
        """Prefix to append the llm call with."""
        return "Thought:"

    @classmethod
    def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
        super()._validate_tools(tools)
        validate_tools_single_input(cls.__name__, tools)

    @classmethod
    def create_prompt(
        cls,
        tools: Sequence[BaseTool],
        system_message: str = PREFIX,
        human_message: str = SUFFIX,
        input_variables: Optional[List[str]] = None,
        output_parser: Optional[BaseOutputParser] = None,
    ) -> BasePromptTemplate:
        # tools
        tool_strings = "\n".join([f"> {tool.name}: {tool.description}" for tool in tools])
        tool_names = ", ".join([tool.name for tool in tools])
        _output_parser = output_parser or cls._get_default_output_parser()

        # format instructions for system message
        format_instructions = _output_parser.get_format_instructions()
        format_instructions = format_instructions.format(tool_names=tool_names)

        # system message


    @property
    def _agent_type(self) -> str:
        raise NotImplementedError

    @property
    def observation_prefix(self) -> str:
        """Prefix to append the observation with."""
        return "Observation: "

    @property
    def llm_prefix(self) -> str:
        """Prefix to append the llm call with."""
        return "Thought:"

    @classmethod
    def _validate_tools(cls, tools: Sequence[BaseTool]) -> None:
        super()._validate_tools(tools)
        validate_tools_single_input(cls.__name__, tools)

    @classmethod
    def create_prompt(
        cls,
        tools: Sequence[BaseTool],
        system_message: str = PREFIX,
        human_message: str = SUFFIX,
        input_variables: Optional[List[str]] = None,
        output_parser: Optional[BaseOutputParser] = None,
    ) -> BasePromptTemplate:
        # tools
        tool_strings = "\n".join([f"> {tool.name}: {tool.description}" for tool in tools])
        tool_names = ", ".join([tool.name for tool in tools])
        _output_parser = output_parser or cls._get_default_output_parser()

        # format instructions for system message
        format_instructions = _output_parser.get_format_instructions()
        format_instructions = format_instructions.format(tool_names=tool_names)

        # system message
        system_message = system_message + f"{tool_strings}\n\n{format_instructions}"

        # human input
        final_prompt = human_message
        if input_variables is None:
            input_variables = ["input", "chat_history", "agent_scratchpad"]
        messages = [
            SystemMessagePromptTemplate.from_template(system_message),
            MessagesPlaceholder(variable_name="chat_history"),
            HumanMessagePromptTemplate.from_template(final_prompt),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]
        return ChatPromptTemplate(input_variables=input_variables, messages=messages)

    @override
    def _construct_scratchpad(self, intermediate_steps: List[Tuple[AgentAction, str]]) -> List[BaseMessage]:
        """Construct the scratchpad that lets the agent continue its thought process."""
        thoughts: List[BaseMessage] = []

        # Try to only use AI message for scratchpad
        content = []
        for idx, (action, full_observation) in enumerate(intermediate_steps):
            content.append(MessageDataModel.extract_action_for_llm(action.log))

            observation = full_observation
            if isinstance(full_observation, DataModel):
                llm_raw_observation = full_observation.get_llm_side_data()
                observation = MessageDataModel.extract_tool_response_for_llm(llm_raw_observation)
                tool_response = self.template_tool_response.format(
                    observation=str(observation), tool_names=self.allowed_tools
                )
                if idx == len(intermediate_steps) - 1:
                    content.append(tool_response)


        # format instructions for system message
        format_instructions = _output_parser.get_format_instructions()
        format_instructions = format_instructions.format(tool_names=tool_names)

        # system message
        system_message = system_message + f"{tool_strings}\n\n{format_instructions}"

        # human input
        final_prompt = human_message
        if input_variables is None:
            input_variables = ["input", "chat_history", "agent_scratchpad"]
        messages = [
            SystemMessagePromptTemplate.from_template(system_message),
            MessagesPlaceholder(variable_name="chat_history"),
            HumanMessagePromptTemplate.from_template(final_prompt),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]
        return ChatPromptTemplate(input_variables=input_variables, messages=messages)

    @override
    def _construct_scratchpad(self, intermediate_steps: List[Tuple[AgentAction, str]]) -> List[BaseMessage]:
        """Construct the scratchpad that lets the agent continue its thought process."""
        thoughts: List[BaseMessage] = []

        # Try to only use AI message for scratchpad
        content = []
        for idx, (action, full_observation) in enumerate(intermediate_steps):
            content.append(MessageDataModel.extract_action_for_llm(action.log))

            observation = full_observation
            if isinstance(full_observation, DataModel):
                llm_raw_observation = full_observation.get_llm_side_data()
                observation = MessageDataModel.extract_tool_response_for_llm(llm_raw_observation)
                tool_response = self.template_tool_response.format(
                    observation=str(observation), tool_names=self.allowed_tools
                )
                if idx == len(intermediate_steps) - 1:
                    content.append(tool_response)
                else:
                    content.append(observation)
        content_str = "\n".join(content)
        thoughts.append(AIMessage(content=content_str))
        if self.continue_model is not None and len(intermediate_steps) != 0:
            thoughts.append(HumanMessage(content=fake_continue_prompt[self.continue_model]))
        return thoughts

    @override
    def plan(
        self,
        intermediate_steps: List[Tuple[AgentAction, str]],
        callbacks: Callbacks = None,
        **kwargs: Any,
    ) -> Union[AgentAction, AgentFinish]:
        """Given input, decided what to do.

        Args:
            intermediate_steps: Steps the LLM has taken to date,
                along with observations
            callbacks: Callbacks to run.
            **kwargs: User inputs.

        Returns:
            Action specifying what tool to use.
        """
        full_inputs = self.get_full_inputs(intermediate_steps, **kwargs)


plan

            content.append(MessageDataModel.extract_action_for_llm(action.log))

            observation = full_observation
            if isinstance(full_observation, DataModel):
                llm_raw_observation = full_observation.get_llm_side_data()
                observation = MessageDataModel.extract_tool_response_for_llm(llm_raw_observation)
                tool_response = self.template_tool_response.format(
                    observation=str(observation), tool_names=self.allowed_tools
                )
                if idx == len(intermediate_steps) - 1:
                    content.append(tool_response)
                else:
                    content.append(observation)
        content_str = "\n".join(content)
        thoughts.append(AIMessage(content=content_str))
        if self.continue_model is not None and len(intermediate_steps) != 0:
            thoughts.append(HumanMessage(content=fake_continue_prompt[self.continue_model]))
        return thoughts

    @override
    def plan(
        self,
        intermediate_steps: List[Tuple[AgentAction, str]],
        callbacks: Callbacks = None,
        **kwargs: Any,
    ) -> Union[AgentAction, AgentFinish]:
        """Given input, decided what to do.

        Args:
            intermediate_steps: Steps the LLM has taken to date,
                along with observations
            callbacks: Callbacks to run.
            **kwargs: User inputs.

        Returns:
            Action specifying what tool to use.
        """
        full_inputs = self.get_full_inputs(intermediate_steps, **kwargs)
        system_prompt = self.llm_chain.prompt.messages[0].format().content
        system_prompt_tokens = MessageDataModel._count_tokens(system_prompt)
        max_tokens = 8000
        max_gen_tokens = 1000
        # FIXME: need more accurate token limit calculation
        full_inputs = MessageDataModel.truncate_chat_history(
            full_inputs, max_token=max_tokens - system_prompt_tokens - max_gen_tokens
        )
        full_output = self.llm_chain.predict(callbacks=callbacks, **full_inputs)

        return self.output_parser.parse(full_output)

    @classmethod
    def from_llm_and_tools(
        cls,
        llm: BaseLanguageModel,
        tools: Sequence[BaseTool],
        callbacks: Callbacks = None,
        output_parser: Optional[AgentOutputParser] = None,
        system_message: str = PREFIX,
        human_message: str = SUFFIX,
        input_variables: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Agent:
        """Construct an agent from an LLM and tools."""
        cls._validate_tools(tools)

        _output_parser = output_parser or cls._get_default_output_parser()
        prompt = cls.create_prompt(
            tools,
            system_message=system_message,


            callbacks: Callbacks to run.
            **kwargs: User inputs.

        Returns:
            Action specifying what tool to use.
        """
        full_inputs = self.get_full_inputs(intermediate_steps, **kwargs)
        system_prompt = self.llm_chain.prompt.messages[0].format().content
        system_prompt_tokens = MessageDataModel._count_tokens(system_prompt)
        max_tokens = 8000
        max_gen_tokens = 1000
        # FIXME: need more accurate token limit calculation
        full_inputs = MessageDataModel.truncate_chat_history(
            full_inputs, max_token=max_tokens - system_prompt_tokens - max_gen_tokens
        )
        full_output = self.llm_chain.predict(callbacks=callbacks, **full_inputs)

        return self.output_parser.parse(full_output)

    @classmethod
    def from_llm_and_tools(
        cls,
        llm: BaseLanguageModel,
        tools: Sequence[BaseTool],
        callbacks: Callbacks = None,
        output_parser: Optional[AgentOutputParser] = None,
        system_message: str = PREFIX,
        human_message: str = SUFFIX,
        input_variables: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Agent:
        """Construct an agent from an LLM and tools."""
        cls._validate_tools(tools)

        _output_parser = output_parser or cls._get_default_output_parser()
        prompt = cls.create_prompt(
            tools,
            system_message=system_message,
            human_message=human_message,
            input_variables=input_variables,
            output_parser=_output_parser,
        )
        llm_chain = LLMChain(
            llm=llm,
            prompt=prompt,
        )
        tool_names = [tool.name for tool in tools]
        return cls(
            llm_chain=llm_chain,
            allowed_tools=tool_names,
            output_parser=_output_parser,
            **kwargs,
        )


import os
from typing import Any, List, Optional, Tuple, Dict
from pydantic import BaseModel
import requests
import time
import ast

import pandas as pd
from io import StringIO
import redis
from loguru import logger

from IPython.core.interactiveshell import InteractiveShell
from IPython.core.getipython import get_ipython
from IPython.utils.capture import capture_output

# subscribed channels
SUBMIT_EVENT = "job_submitted"
RUNNING_EVENT = "job_started"
COMPLETE_EVENT = "job_completed"
# Error render prefix
ERROR_PREFIX = "[ERROR]: "

def check_danger_code(code):
    code_line = []
    for line in code.split("\n"):
        if not line.startswith("%"):
            code_line.append(line)
    code = "\n".join(code_line)

    def check_imports(code):
        ast_failed = False
        try:
            tree = ast.parse(code)
        except Exception as e:
            ast_failed = str(e)
            return ast_failed
        return ast_failed

    ast_failed = check_imports(code)
    return True, ast_failed, []

class DisplayData(BaseModel):
    """Both display_data and execute_result messages use this format."""

    data: Optional[dict] = None
    metadata: Optional[dict] = None

    @classmethod
    def from_tuple(cls, formatted: Tuple[dict, dict]):
        return cls(data=formatted[0], metadata=formatted[1])

    def to_dict(self) -> Dict:
        return {
            "data": self.data,
            "metadata": self.metadata,
        }

class PythonEvaluator:
    """
    Util class for Python code evaluation.
    """

    name = "Python Evaluator"
    base_url = "http://{0}:8100".format(os.getenv("CODE_INTER_SERVER"))
    r: redis.Redis = redis.Redis(host=os.getenv("REDIS_SERVER"), port=6379, decode_responses=True)

    def __init__(self, code_execution_mode: str = "local", jupyter_kernel_pool: Optional[Any] = None):
        self.code_execution_mode = code_execution_mode
        self.jupyter_kernel_pool = jupyter_kernel_pool

    @staticmethod
    def parse_command(program: str) -> List[str]:
        """patchify the code"""
        program_lines = program.strip().split("\n")
        return program_lines

    def run_program_local(self, program: str, user_id: Optional[str] = "u" * 24):
        """Run python program on the local machine using Ipython shell."""
        is_safe, ast_failed, danger_pcks = check_danger_code(program)
        if ast_failed != False:
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}Error Code Parsing, please check code grammar!\n{ast_failed}",
            }
        try:
            # Run code using local ipython shell
            # Change working dir to data directory to load from pretty path
            #   Note! This is not thread safe, only for local use
            os.chdir(os.path.join("backend/data/", user_id))

            shell = InteractiveShell.instance()
            shell.enable_gui = lambda x: False
            with capture_output() as captured:
                ip = get_ipython()
                code = "%matplotlib inline\n" + program  # magic command to display matplotlib plots
                result = ip.run_cell(code)

            # Change working dir to project root
            os.chdir("../../../")

            if result.success:
                return {
                    "success": True,
                    "result": result.result,
                    "stdout": str(captured.stdout),
                    "stderr": str(captured.stderr),
                    "outputs": captured.outputs,
                }
            elif result.error_in_exec is not None:
                return {
                    "success": False,
                    "error_message": f"{ERROR_PREFIX}{str(result.error_in_exec)}",
                    "outputs": captured.outputs,
                }
            else:
                # error_before_exec
                return {
                    "success": False,
                    "error_message": f"{ERROR_PREFIX}{str(result.error_before_exec)}",
                    "outputs": captured.outputs,
                }
        except Exception as e:
            logger.bind(user_id=user_id, msg_head="Python evaluator running error").trace(e)
            import traceback

            traceback.print_exc()
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}{str(e)}",
            }

    def _apply_for_kernel(self, kernel_id: Optional[str], user_id: str, chat_id: str):
        """Apply for a kernel in docker to run program."""
        if kernel_id is not None:
            # If kernel id is provided, use it directly
            cur_kid = kernel_id
        else:
            # If kernel id is not provided, apply for a new kernel
            kernel_info = self.jupyter_kernel_pool.get_pool_info_with_id(user_id, chat_id, None)
            cur_kid = kernel_info["kid"] if kernel_info is not None else None
            user_exists = requests.get(f"{self.base_url}/user/status/{user_id}").json()["exists"]

            logger.bind(user_id=user_id, chat_id=chat_id, msg_head="user exists").trace(user_exists)

            if not user_exists:
                response = requests.post(f"{self.base_url}/user/create", json={"username": user_id}).json()

                logger.bind(user_id=user_id, chat_id=chat_id, msg_head="user create").trace(response)

            response = requests.get(f"{self.base_url}/kernel/list/{user_id}").json()
            existing_kernel_list = response["list"]

            logger.bind(user_id=user_id, chat_id=chat_id, msg_head="kernel list").trace(response)

            if cur_kid not in existing_kernel_list:
                response = requests.post(f"{self.base_url}/kernel/create", json={"username": user_id}).json()
                if response["code"] != 0 and response["msg"] == "Too many kernels":
                    # kill oldest kernel
                    oldest_kernel_id = existing_kernel_list[0]
                    response = requests.post(
                        f"{self.base_url}/kernel/stop", json={"username": user_id, "kid": oldest_kernel_id}
                    ).json()

                    logger.bind(user_id=user_id, chat_id=chat_id, msg_head="kill oldest kernel").trace(response)

                    response = requests.post(f"{self.base_url}/kernel/create", json={"username": user_id}).json()
                cur_kid = response["id"]

                logger.bind(user_id=user_id, chat_id=chat_id, msg_head="create kernel id").trace(cur_kid)

                self.jupyter_kernel_pool.set_pool_info_with_id(
                    user_id, chat_id, {"kid": cur_kid, "ktime": time.time()}
                )

        logger.bind(user_id=user_id, chat_id=chat_id, msg_head="current kernel id").trace(cur_kid)

        return cur_kid

    def run_program_docker(
        self,
        program: str,
        kernel_id: Optional[str] = None,
        user_id: Optional[str] = "u" * 24,
        chat_id: Optional[str] = "c" * 24,
    ):
        """Run python program on the docker container(jupyter client)."""
        is_safe, ast_failed, danger_pcks = check_danger_code(program)
        if not is_safe:
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}Unsafe Code Detected {str(danger_pcks)}, Execution Denied!",
            }
        elif ast_failed != False:
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}Error Code Parsing, please check code grammar!\n{ast_failed}",
            }

        try:
            # Run code using remote docker jupyter kernel
            # Get the running id(i.e., a permission to run program) from redis running queue
            #  otherwise wait until a running_id is available
            p = self.r.pubsub()
            p.subscribe(RUNNING_EVENT)
            self.r.publish(SUBMIT_EVENT, chat_id)
            for message in p.listen():
                # the initial message for each channel is a message with an integer
                if isinstance(message["data"], int):
                    continue
                running_id = message["data"]
                if running_id == chat_id:
                    break
                time.sleep(1)
            # Get kernel id(i.e., the real jupyter kernel to run the program) to execute program
            cur_kid = self._apply_for_kernel(kernel_id, user_id, chat_id)
            # Execute program
            response = requests.post(
                f"{self.base_url}/kernel/exec", json={"username": user_id, "code": program, "kid": cur_kid}
            ).json()
            # Notify Redis that a job has been completed
            self.r.publish(COMPLETE_EVENT, chat_id)

            # Parse jupyter kernel output
            result, stdout, stderr, outputs, displays, error_message = None, "", "", None, [], None
            if response["status"] == "ok":
                output = response.get("output", None)
                if output is not None:
                    for output_dict in output:
                        if output_dict["type"] == "stream":
                            content = output_dict.get("content", None)
                            if content is not None:
                                if content["name"] == "stdout":
                                    stdout = content["text"]
                                if content["name"] == "stderr":
                                    stderr = content["text"]
                        elif output_dict["type"] == "execute_result":
                            content = output_dict.get("content", None)
                            if content is not None:
                                data = content.get("data", None)
                                if data is not None:
                                    if "text/plain" in data and "text/html" in data:
                                        try:
                                            # Try to recover a dataframe
                                            df = pd.read_csv(StringIO(data["text/plain"]))
                                            result = df
                                        except Exception as e:
                                            pass
                                    elif "text/plain" in data:
                                        result = data["text/plain"]
                                    else:
                                        # TODO: If not match any of the above, return the first value
                                        result = list(data.values())[0]
                        elif output_dict["type"] == "display_data":
                            content = output_dict.get("content", None)
                            if content is not None:
                                data = content.get("data", None)
                                metadata = content.get("metadata", None)
                                if data is not None and metadata is not None:
                                    if "image/png" in data:
                                        displays.append(DisplayData.from_tuple((data, metadata)))
                            if displays:
                                outputs = displays
            else:
                return {
                    "success": False,
                    "error_message": f"{ERROR_PREFIX}Error status returned by kernel",
                }
            # Check success status and return
            shell_msg = response["shell"]
            if shell_msg["status"] == "ok":
                return {
                    "success": True,
                    "result": result,
                    "stdout": stdout,
                    "stderr": stderr,
                    "outputs": outputs,
                }
            elif shell_msg["status"] == "error":
                error_message = f"{shell_msg['ename']}: {shell_msg['evalue']}"
                return {"success": False, "error_message": f"{ERROR_PREFIX}{error_message}", "outputs": outputs}
        except Exception as e:
            logger.bind(user_id=user_id, chat_id=chat_id, msg_head="Python evaluator running error").trace(e)
            import traceback

            traceback.print_exc()

            try:
                # Notify Redis that a job has been completed
                self.r.publish(COMPLETE_EVENT, chat_id)
            except:
                pass

            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}{str(e)}",
            }

    def run(
        self,
        program: str,
        environment: Optional[Any] = None,
        kernel_id: Optional[str] = None,
        user_id: Optional[str] = "u" * 24,
        chat_id: Optional[str] = "c" * 24,
    ) -> Any:
        """run generated code in certain environment"""

        lines_code = self.parse_command(program)
        program = "\n".join(lines_code)
        program = "%matplotlib inline\n" + program  # magic command to display matplotlib plots

        logger.bind(user_id=user_id, chat_id=chat_id, msg_head="Code execution mode").trace(self.code_execution_mode)

        if self.code_execution_mode == "local":
            return self.run_program_local(program, user_id)
        elif self.code_execution_mode == "docker":
            return self.run_program_docker(program, kernel_id, user_id, chat_id)
        else:
            raise ValueError("Invalid code execution mode")


import ast

import pandas as pd
from io import StringIO
import redis
from loguru import logger

from IPython.core.interactiveshell import InteractiveShell
from IPython.core.getipython import get_ipython
from IPython.utils.capture import capture_output

# subscribed channels
SUBMIT_EVENT = "job_submitted"
RUNNING_EVENT = "job_started"
COMPLETE_EVENT = "job_completed"
# Error render prefix
ERROR_PREFIX = "[ERROR]: "

def check_danger_code(code):
    code_line = []
    for line in code.split("\n"):
        if not line.startswith("%"):
            code_line.append(line)
    code = "\n".join(code_line)

    def check_imports(code):
        ast_failed = False
        try:
            tree = ast.parse(code)
        except Exception as e:
            ast_failed = str(e)
            return ast_failed
        return ast_failed

    ast_failed = check_imports(code)
    return True, ast_failed, []

class DisplayData(BaseModel):
    """Both display_data and execute_result messages use this format."""

    data: Optional[dict] = None
    metadata: Optional[dict] = None

    @classmethod
    def from_tuple(cls, formatted: Tuple[dict, dict]):
        return cls(data=formatted[0], metadata=formatted[1])

    def to_dict(self) -> Dict:
        return {
            "data": self.data,
            "metadata": self.metadata,
        }

class PythonEvaluator:


DisplayData

def check_danger_code(code):
    code_line = []
    for line in code.split("\n"):
        if not line.startswith("%"):
            code_line.append(line)
    code = "\n".join(code_line)

    def check_imports(code):
        ast_failed = False
        try:
            tree = ast.parse(code)
        except Exception as e:
            ast_failed = str(e)
            return ast_failed
        return ast_failed

    ast_failed = check_imports(code)
    return True, ast_failed, []

class DisplayData(BaseModel):
    """Both display_data and execute_result messages use this format."""

    data: Optional[dict] = None
    metadata: Optional[dict] = None

    @classmethod
    def from_tuple(cls, formatted: Tuple[dict, dict]):
        return cls(data=formatted[0], metadata=formatted[1])

    def to_dict(self) -> Dict:
        return {
            "data": self.data,
            "metadata": self.metadata,
        }

class PythonEvaluator:
    """
    Util class for Python code evaluation.
    """

    name = "Python Evaluator"
    base_url = "http://{0}:8100".format(os.getenv("CODE_INTER_SERVER"))
    r: redis.Redis = redis.Redis(host=os.getenv("REDIS_SERVER"), port=6379, decode_responses=True)

    def __init__(self, code_execution_mode: str = "local", jupyter_kernel_pool: Optional[Any] = None):
        self.code_execution_mode = code_execution_mode
        self.jupyter_kernel_pool = jupyter_kernel_pool

    @staticmethod
    def parse_command(program: str) -> List[str]:
        """patchify the code"""
        program_lines = program.strip().split("\n")
        return program_lines


PythonEvaluator

    return True, ast_failed, []

class DisplayData(BaseModel):
    """Both display_data and execute_result messages use this format."""

    data: Optional[dict] = None
    metadata: Optional[dict] = None

    @classmethod
    def from_tuple(cls, formatted: Tuple[dict, dict]):
        return cls(data=formatted[0], metadata=formatted[1])

    def to_dict(self) -> Dict:
        return {
            "data": self.data,
            "metadata": self.metadata,
        }

class PythonEvaluator:
    """
    Util class for Python code evaluation.
    """

    name = "Python Evaluator"
    base_url = "http://{0}:8100".format(os.getenv("CODE_INTER_SERVER"))
    r: redis.Redis = redis.Redis(host=os.getenv("REDIS_SERVER"), port=6379, decode_responses=True)

    def __init__(self, code_execution_mode: str = "local", jupyter_kernel_pool: Optional[Any] = None):
        self.code_execution_mode = code_execution_mode
        self.jupyter_kernel_pool = jupyter_kernel_pool

    @staticmethod
    def parse_command(program: str) -> List[str]:
        """patchify the code"""
        program_lines = program.strip().split("\n")
        return program_lines

    def run_program_local(self, program: str, user_id: Optional[str] = "u" * 24):
        """Run python program on the local machine using Ipython shell."""
        is_safe, ast_failed, danger_pcks = check_danger_code(program)
        if ast_failed != False:
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}Error Code Parsing, please check code grammar!\n{ast_failed}",
            }
        try:
            # Run code using local ipython shell
            # Change working dir to data directory to load from pretty path
            #   Note! This is not thread safe, only for local use
            os.chdir(os.path.join("backend/data/", user_id))

            shell = InteractiveShell.instance()
            shell.enable_gui = lambda x: False
            with capture_output() as captured:
                ip = get_ipython()
                code = "%matplotlib inline\n" + program  # magic command to display matplotlib plots
                result = ip.run_cell(code)

            # Change working dir to project root
            os.chdir("../../../")

            if result.success:
                return {
                    "success": True,
                    "result": result.result,
                    "stdout": str(captured.stdout),
                    "stderr": str(captured.stderr),
                    "outputs": captured.outputs,
                }
            elif result.error_in_exec is not None:
                return {
                    "success": False,
                    "error_message": f"{ERROR_PREFIX}{str(result.error_in_exec)}",
                    "outputs": captured.outputs,
                }
            else:
                # error_before_exec
                return {
                    "success": False,
                    "error_message": f"{ERROR_PREFIX}{str(result.error_before_exec)}",
                    "outputs": captured.outputs,
                }
        except Exception as e:
            logger.bind(user_id=user_id, msg_head="Python evaluator running error").trace(e)
            import traceback

            traceback.print_exc()
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}{str(e)}",
            }

    def _apply_for_kernel(self, kernel_id: Optional[str], user_id: str, chat_id: str):
        """Apply for a kernel in docker to run program."""
        if kernel_id is not None:
            # If kernel id is provided, use it directly
            cur_kid = kernel_id
        else:
            # If kernel id is not provided, apply for a new kernel
            kernel_info = self.jupyter_kernel_pool.get_pool_info_with_id(user_id, chat_id, None)
            cur_kid = kernel_info["kid"] if kernel_info is not None else None
            user_exists = requests.get(f"{self.base_url}/user/status/{user_id}").json()["exists"]

            logger.bind(user_id=user_id, chat_id=chat_id, msg_head="user exists").trace(user_exists)

            if not user_exists:
                response = requests.post(f"{self.base_url}/user/create", json={"username": user_id}).json()

                logger.bind(user_id=user_id, chat_id=chat_id, msg_head="user create").trace(response)

            response = requests.get(f"{self.base_url}/kernel/list/{user_id}").json()
            existing_kernel_list = response["list"]

            logger.bind(user_id=user_id, chat_id=chat_id, msg_head="kernel list").trace(response)

            if cur_kid not in existing_kernel_list:
                response = requests.post(f"{self.base_url}/kernel/create", json={"username": user_id}).json()
                if response["code"] != 0 and response["msg"] == "Too many kernels":
                    # kill oldest kernel
                    oldest_kernel_id = existing_kernel_list[0]
                    response = requests.post(
                        f"{self.base_url}/kernel/stop", json={"username": user_id, "kid": oldest_kernel_id}
                    ).json()

                    logger.bind(user_id=user_id, chat_id=chat_id, msg_head="kill oldest kernel").trace(response)

                    response = requests.post(f"{self.base_url}/kernel/create", json={"username": user_id}).json()
                cur_kid = response["id"]

                logger.bind(user_id=user_id, chat_id=chat_id, msg_head="create kernel id").trace(cur_kid)

                self.jupyter_kernel_pool.set_pool_info_with_id(
                    user_id, chat_id, {"kid": cur_kid, "ktime": time.time()}
                )

        logger.bind(user_id=user_id, chat_id=chat_id, msg_head="current kernel id").trace(cur_kid)

        return cur_kid

    def run_program_docker(
        self,
        program: str,
        kernel_id: Optional[str] = None,
        user_id: Optional[str] = "u" * 24,
        chat_id: Optional[str] = "c" * 24,
    ):
        """Run python program on the docker container(jupyter client)."""
        is_safe, ast_failed, danger_pcks = check_danger_code(program)
        if not is_safe:
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}Unsafe Code Detected {str(danger_pcks)}, Execution Denied!",
            }
        elif ast_failed != False:
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}Error Code Parsing, please check code grammar!\n{ast_failed}",
            }

        try:
            # Run code using remote docker jupyter kernel
            # Get the running id(i.e., a permission to run program) from redis running queue
            #  otherwise wait until a running_id is available
            p = self.r.pubsub()
            p.subscribe(RUNNING_EVENT)
            self.r.publish(SUBMIT_EVENT, chat_id)
            for message in p.listen():
                # the initial message for each channel is a message with an integer
                if isinstance(message["data"], int):
                    continue
                running_id = message["data"]
                if running_id == chat_id:
                    break
                time.sleep(1)
            # Get kernel id(i.e., the real jupyter kernel to run the program) to execute program
            cur_kid = self._apply_for_kernel(kernel_id, user_id, chat_id)
            # Execute program
            response = requests.post(
                f"{self.base_url}/kernel/exec", json={"username": user_id, "code": program, "kid": cur_kid}
            ).json()
            # Notify Redis that a job has been completed
            self.r.publish(COMPLETE_EVENT, chat_id)

            # Parse jupyter kernel output
            result, stdout, stderr, outputs, displays, error_message = None, "", "", None, [], None
            if response["status"] == "ok":
                output = response.get("output", None)
                if output is not None:
                    for output_dict in output:
                        if output_dict["type"] == "stream":
                            content = output_dict.get("content", None)
                            if content is not None:
                                if content["name"] == "stdout":
                                    stdout = content["text"]
                                if content["name"] == "stderr":
                                    stderr = content["text"]
                        elif output_dict["type"] == "execute_result":
                            content = output_dict.get("content", None)
                            if content is not None:
                                data = content.get("data", None)
                                if data is not None:
                                    if "text/plain" in data and "text/html" in data:
                                        try:
                                            # Try to recover a dataframe
                                            df = pd.read_csv(StringIO(data["text/plain"]))
                                            result = df
                                        except Exception as e:
                                            pass
                                    elif "text/plain" in data:
                                        result = data["text/plain"]
                                    else:
                                        # TODO: If not match any of the above, return the first value
                                        result = list(data.values())[0]
                        elif output_dict["type"] == "display_data":
                            content = output_dict.get("content", None)
                            if content is not None:
                                data = content.get("data", None)
                                metadata = content.get("metadata", None)
                                if data is not None and metadata is not None:
                                    if "image/png" in data:
                                        displays.append(DisplayData.from_tuple((data, metadata)))
                            if displays:
                                outputs = displays
            else:
                return {
                    "success": False,
                    "error_message": f"{ERROR_PREFIX}Error status returned by kernel",
                }
            # Check success status and return
            shell_msg = response["shell"]
            if shell_msg["status"] == "ok":
                return {
                    "success": True,
                    "result": result,
                    "stdout": stdout,
                    "stderr": stderr,
                    "outputs": outputs,
                }
            elif shell_msg["status"] == "error":
                error_message = f"{shell_msg['ename']}: {shell_msg['evalue']}"
                return {"success": False, "error_message": f"{ERROR_PREFIX}{error_message}", "outputs": outputs}
        except Exception as e:
            logger.bind(user_id=user_id, chat_id=chat_id, msg_head="Python evaluator running error").trace(e)
            import traceback

            traceback.print_exc()

            try:
                # Notify Redis that a job has been completed
                self.r.publish(COMPLETE_EVENT, chat_id)
            except:
                pass

            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}{str(e)}",
            }

    def run(
        self,
        program: str,
        environment: Optional[Any] = None,
        kernel_id: Optional[str] = None,
        user_id: Optional[str] = "u" * 24,
        chat_id: Optional[str] = "c" * 24,
    ) -> Any:
        """run generated code in certain environment"""

        lines_code = self.parse_command(program)
        program = "\n".join(lines_code)
        program = "%matplotlib inline\n" + program  # magic command to display matplotlib plots

        logger.bind(user_id=user_id, chat_id=chat_id, msg_head="Code execution mode").trace(self.code_execution_mode)

        if self.code_execution_mode == "local":
            return self.run_program_local(program, user_id)
        elif self.code_execution_mode == "docker":
            return self.run_program_docker(program, kernel_id, user_id, chat_id)
        else:
            raise ValueError("Invalid code execution mode")


from IPython.core.interactiveshell import InteractiveShell
from IPython.core.getipython import get_ipython
from IPython.utils.capture import capture_output

# subscribed channels
SUBMIT_EVENT = "job_submitted"
RUNNING_EVENT = "job_started"
COMPLETE_EVENT = "job_completed"
# Error render prefix
ERROR_PREFIX = "[ERROR]: "

def check_danger_code(code):
    code_line = []
    for line in code.split("\n"):
        if not line.startswith("%"):
            code_line.append(line)
    code = "\n".join(code_line)

    def check_imports(code):
        ast_failed = False
        try:
            tree = ast.parse(code)
        except Exception as e:
            ast_failed = str(e)
            return ast_failed
        return ast_failed

    ast_failed = check_imports(code)
    return True, ast_failed, []

class DisplayData(BaseModel):
    """Both display_data and execute_result messages use this format."""

    data: Optional[dict] = None
    metadata: Optional[dict] = None

    @classmethod
    def from_tuple(cls, formatted: Tuple[dict, dict]):
        return cls(data=formatted[0], metadata=formatted[1])

    def to_dict(self) -> Dict:
        return {
            "data": self.data,
            "metadata": self.metadata,
        }


    def check_imports(code):
        ast_failed = False
        try:
            tree = ast.parse(code)
        except Exception as e:
            ast_failed = str(e)
            return ast_failed
        return ast_failed

    ast_failed = check_imports(code)
    return True, ast_failed, []

class DisplayData(BaseModel):
    """Both display_data and execute_result messages use this format."""

    data: Optional[dict] = None
    metadata: Optional[dict] = None

    @classmethod
    def from_tuple(cls, formatted: Tuple[dict, dict]):
        return cls(data=formatted[0], metadata=formatted[1])

    def to_dict(self) -> Dict:
        return {
            "data": self.data,
            "metadata": self.metadata,
        }

class PythonEvaluator:
    """
    Util class for Python code evaluation.
    """

    name = "Python Evaluator"
    base_url = "http://{0}:8100".format(os.getenv("CODE_INTER_SERVER"))
    r: redis.Redis = redis.Redis(host=os.getenv("REDIS_SERVER"), port=6379, decode_responses=True)

    def __init__(self, code_execution_mode: str = "local", jupyter_kernel_pool: Optional[Any] = None):
        self.code_execution_mode = code_execution_mode
        self.jupyter_kernel_pool = jupyter_kernel_pool


            tree = ast.parse(code)
        except Exception as e:
            ast_failed = str(e)
            return ast_failed
        return ast_failed

    ast_failed = check_imports(code)
    return True, ast_failed, []

class DisplayData(BaseModel):
    """Both display_data and execute_result messages use this format."""

    data: Optional[dict] = None
    metadata: Optional[dict] = None

    @classmethod
    def from_tuple(cls, formatted: Tuple[dict, dict]):
        return cls(data=formatted[0], metadata=formatted[1])

    def to_dict(self) -> Dict:
        return {
            "data": self.data,
            "metadata": self.metadata,
        }

class PythonEvaluator:
    """
    Util class for Python code evaluation.
    """

    name = "Python Evaluator"
    base_url = "http://{0}:8100".format(os.getenv("CODE_INTER_SERVER"))
    r: redis.Redis = redis.Redis(host=os.getenv("REDIS_SERVER"), port=6379, decode_responses=True)

    def __init__(self, code_execution_mode: str = "local", jupyter_kernel_pool: Optional[Any] = None):
        self.code_execution_mode = code_execution_mode
        self.jupyter_kernel_pool = jupyter_kernel_pool

    @staticmethod
    def parse_command(program: str) -> List[str]:
        """patchify the code"""
        program_lines = program.strip().split("\n")
        return program_lines


    @classmethod
    def from_tuple(cls, formatted: Tuple[dict, dict]):
        return cls(data=formatted[0], metadata=formatted[1])

    def to_dict(self) -> Dict:
        return {
            "data": self.data,
            "metadata": self.metadata,
        }

class PythonEvaluator:
    """
    Util class for Python code evaluation.
    """

    name = "Python Evaluator"
    base_url = "http://{0}:8100".format(os.getenv("CODE_INTER_SERVER"))
    r: redis.Redis = redis.Redis(host=os.getenv("REDIS_SERVER"), port=6379, decode_responses=True)

    def __init__(self, code_execution_mode: str = "local", jupyter_kernel_pool: Optional[Any] = None):
        self.code_execution_mode = code_execution_mode
        self.jupyter_kernel_pool = jupyter_kernel_pool

    @staticmethod
    def parse_command(program: str) -> List[str]:
        """patchify the code"""
        program_lines = program.strip().split("\n")
        return program_lines

    def run_program_local(self, program: str, user_id: Optional[str] = "u" * 24):
        """Run python program on the local machine using Ipython shell."""
        is_safe, ast_failed, danger_pcks = check_danger_code(program)
        if ast_failed != False:
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}Error Code Parsing, please check code grammar!\n{ast_failed}",
            }
        try:
            # Run code using local ipython shell
            # Change working dir to data directory to load from pretty path
            #   Note! This is not thread safe, only for local use
            os.chdir(os.path.join("backend/data/", user_id))


        return {
            "data": self.data,
            "metadata": self.metadata,
        }

class PythonEvaluator:
    """
    Util class for Python code evaluation.
    """

    name = "Python Evaluator"
    base_url = "http://{0}:8100".format(os.getenv("CODE_INTER_SERVER"))
    r: redis.Redis = redis.Redis(host=os.getenv("REDIS_SERVER"), port=6379, decode_responses=True)

    def __init__(self, code_execution_mode: str = "local", jupyter_kernel_pool: Optional[Any] = None):
        self.code_execution_mode = code_execution_mode
        self.jupyter_kernel_pool = jupyter_kernel_pool

    @staticmethod
    def parse_command(program: str) -> List[str]:
        """patchify the code"""
        program_lines = program.strip().split("\n")
        return program_lines

    def run_program_local(self, program: str, user_id: Optional[str] = "u" * 24):
        """Run python program on the local machine using Ipython shell."""
        is_safe, ast_failed, danger_pcks = check_danger_code(program)
        if ast_failed != False:
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}Error Code Parsing, please check code grammar!\n{ast_failed}",
            }
        try:
            # Run code using local ipython shell
            # Change working dir to data directory to load from pretty path
            #   Note! This is not thread safe, only for local use
            os.chdir(os.path.join("backend/data/", user_id))

            shell = InteractiveShell.instance()
            shell.enable_gui = lambda x: False
            with capture_output() as captured:
                ip = get_ipython()
                code = "%matplotlib inline\n" + program  # magic command to display matplotlib plots


class PythonEvaluator:
    """
    Util class for Python code evaluation.
    """

    name = "Python Evaluator"
    base_url = "http://{0}:8100".format(os.getenv("CODE_INTER_SERVER"))
    r: redis.Redis = redis.Redis(host=os.getenv("REDIS_SERVER"), port=6379, decode_responses=True)

    def __init__(self, code_execution_mode: str = "local", jupyter_kernel_pool: Optional[Any] = None):
        self.code_execution_mode = code_execution_mode
        self.jupyter_kernel_pool = jupyter_kernel_pool

    @staticmethod
    def parse_command(program: str) -> List[str]:
        """patchify the code"""
        program_lines = program.strip().split("\n")
        return program_lines

    def run_program_local(self, program: str, user_id: Optional[str] = "u" * 24):
        """Run python program on the local machine using Ipython shell."""
        is_safe, ast_failed, danger_pcks = check_danger_code(program)
        if ast_failed != False:
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}Error Code Parsing, please check code grammar!\n{ast_failed}",
            }
        try:
            # Run code using local ipython shell
            # Change working dir to data directory to load from pretty path
            #   Note! This is not thread safe, only for local use
            os.chdir(os.path.join("backend/data/", user_id))

            shell = InteractiveShell.instance()
            shell.enable_gui = lambda x: False
            with capture_output() as captured:
                ip = get_ipython()
                code = "%matplotlib inline\n" + program  # magic command to display matplotlib plots
                result = ip.run_cell(code)

            # Change working dir to project root
            os.chdir("../../../")

            if result.success:
                return {
                    "success": True,
                    "result": result.result,
                    "stdout": str(captured.stdout),
                    "stderr": str(captured.stderr),
                    "outputs": captured.outputs,
                }
            elif result.error_in_exec is not None:
                return {
                    "success": False,
                    "error_message": f"{ERROR_PREFIX}{str(result.error_in_exec)}",
                    "outputs": captured.outputs,
                }
            else:
                # error_before_exec
                return {
                    "success": False,
                    "error_message": f"{ERROR_PREFIX}{str(result.error_before_exec)}",
                    "outputs": captured.outputs,
                }
        except Exception as e:
            logger.bind(user_id=user_id, msg_head="Python evaluator running error").trace(e)
            import traceback

            traceback.print_exc()
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}{str(e)}",
            }

    def _apply_for_kernel(self, kernel_id: Optional[str], user_id: str, chat_id: str):
        """Apply for a kernel in docker to run program."""
        if kernel_id is not None:
            # If kernel id is provided, use it directly
            cur_kid = kernel_id
        else:
            # If kernel id is not provided, apply for a new kernel
            kernel_info = self.jupyter_kernel_pool.get_pool_info_with_id(user_id, chat_id, None)
            cur_kid = kernel_info["kid"] if kernel_info is not None else None
            user_exists = requests.get(f"{self.base_url}/user/status/{user_id}").json()["exists"]

            logger.bind(user_id=user_id, chat_id=chat_id, msg_head="user exists").trace(user_exists)

            if not user_exists:
                response = requests.post(f"{self.base_url}/user/create", json={"username": user_id}).json()

                logger.bind(user_id=user_id, chat_id=chat_id, msg_head="user create").trace(response)

            response = requests.get(f"{self.base_url}/kernel/list/{user_id}").json()


                    "error_message": f"{ERROR_PREFIX}{str(result.error_in_exec)}",
                    "outputs": captured.outputs,
                }
            else:
                # error_before_exec
                return {
                    "success": False,
                    "error_message": f"{ERROR_PREFIX}{str(result.error_before_exec)}",
                    "outputs": captured.outputs,
                }
        except Exception as e:
            logger.bind(user_id=user_id, msg_head="Python evaluator running error").trace(e)
            import traceback

            traceback.print_exc()
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}{str(e)}",
            }

    def _apply_for_kernel(self, kernel_id: Optional[str], user_id: str, chat_id: str):
        """Apply for a kernel in docker to run program."""
        if kernel_id is not None:
            # If kernel id is provided, use it directly
            cur_kid = kernel_id
        else:
            # If kernel id is not provided, apply for a new kernel
            kernel_info = self.jupyter_kernel_pool.get_pool_info_with_id(user_id, chat_id, None)
            cur_kid = kernel_info["kid"] if kernel_info is not None else None
            user_exists = requests.get(f"{self.base_url}/user/status/{user_id}").json()["exists"]

            logger.bind(user_id=user_id, chat_id=chat_id, msg_head="user exists").trace(user_exists)

            if not user_exists:
                response = requests.post(f"{self.base_url}/user/create", json={"username": user_id}).json()

                logger.bind(user_id=user_id, chat_id=chat_id, msg_head="user create").trace(response)

            response = requests.get(f"{self.base_url}/kernel/list/{user_id}").json()
            existing_kernel_list = response["list"]

            logger.bind(user_id=user_id, chat_id=chat_id, msg_head="kernel list").trace(response)

            if cur_kid not in existing_kernel_list:
                response = requests.post(f"{self.base_url}/kernel/create", json={"username": user_id}).json()
                if response["code"] != 0 and response["msg"] == "Too many kernels":
                    # kill oldest kernel
                    oldest_kernel_id = existing_kernel_list[0]
                    response = requests.post(
                        f"{self.base_url}/kernel/stop", json={"username": user_id, "kid": oldest_kernel_id}
                    ).json()

                    logger.bind(user_id=user_id, chat_id=chat_id, msg_head="kill oldest kernel").trace(response)

                    response = requests.post(f"{self.base_url}/kernel/create", json={"username": user_id}).json()
                cur_kid = response["id"]

                logger.bind(user_id=user_id, chat_id=chat_id, msg_head="create kernel id").trace(cur_kid)

                self.jupyter_kernel_pool.set_pool_info_with_id(
                    user_id, chat_id, {"kid": cur_kid, "ktime": time.time()}
                )

        logger.bind(user_id=user_id, chat_id=chat_id, msg_head="current kernel id").trace(cur_kid)

        return cur_kid

    def run_program_docker(
        self,
        program: str,
        kernel_id: Optional[str] = None,
        user_id: Optional[str] = "u" * 24,
        chat_id: Optional[str] = "c" * 24,
    ):
        """Run python program on the docker container(jupyter client)."""
        is_safe, ast_failed, danger_pcks = check_danger_code(program)
        if not is_safe:
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}Unsafe Code Detected {str(danger_pcks)}, Execution Denied!",
            }
        elif ast_failed != False:
            return {
                "success": False,
                "error_message": f"{ERROR_PREFIX}Error Code Parsing, please check code grammar!\n{ast_failed}",
            }