main.py

"""The driver file for the program."""

import re
import os
from typing import Union
from collections import deque
from pathlib import Path

from exec_engine.db_manager import DBManager
from exec_engine.api_executor import APIExecutor, PythonAPIExecutor
from exec_engine.docker_sandbox import DockerSandbox
from exec_engine.fs_manager import FSManager
from exec_engine.utils import SQL_Type, RESTful_Type, Filesystem_Type

from exec_engine.pipeline import generate_command

ROOT_FOLDER_PATH = os.path.dirname(Path(os.path.realpath(__file__)))
CHECKPOINTS_FOLDER_PATH = os.path.join(ROOT_FOLDER_PATH, "docker/mysql_docker/checkpoints")

class ExecutionEngine:
    """Can identify and execute API calls and their reversals."""

    def __init__(self, history_length: int = 10, path=None, generate_mode='default'):
        # Keys for LLMs
        self.openai_api_key = None
        self.anthropic_api_key = None
        # Initialize the instances of the classes
        self.api_executor = APIExecutor()
        self.docker_sandbox = DockerSandbox()
        
        self.fs_initialized = False
        self.db_initialized = False
        self.path = path
        self.generate_mode = generate_mode

        # Dry Run is option 2, no dry run is option 1, user can specify which types to dry-run
        self.dry_run_dict = {
            SQL_Type: False,
            Filesystem_Type: False,
            RESTful_Type: False,
        }

        # Initialize the API and ~API pair circular queue
        self.api_history_pair_queue_dict = {
            SQL_Type: deque(maxlen=history_length),
            Filesystem_Type: deque(maxlen=history_length),
            RESTful_Type: deque(maxlen=history_length),
        }

    def gen_api_pair(self, input_prompt: str, api_type: str, credentials, model) -> list:
        """Generate an API call and its reversal"""
        if api_type == RESTful_Type:
            forward_call = generate_command(input_prompt, credentials, generate_mode=self.generate_mode, openai_model=model)
            backward_call = self.api_executor.try_get_backward_call(forward_call, input_prompt, credentials, api_type,
                                                                     generate_mode=self.generate_mode, model=model)
            return forward_call, backward_call
            
        elif api_type == SQL_Type:
            input_p = self.db_manager.task_to_prompt(input_prompt)
            forward_call = generate_command(input_p, credentials, api_type=SQL_Type, generate_mode=self.generate_mode, openai_model=model)

            reverse_p = self.db_manager.task_to_prompt(forward_call, forward=False)
            backward_call = generate_command(reverse_p, credentials, api_type=SQL_Type, generate_mode=self.generate_mode, openai_model=model)
            return forward_call, backward_call

        elif api_type == Filesystem_Type:
            input_p = self.fs_manager.task_to_prompt(input_prompt)
            forward_call = generate_command(input_p, credentials, api_type=Filesystem_Type, generate_mode=self.generate_mode, openai_model=model)

            reverse_p = self.fs_manager.task_to_prompt(forward_call, forward=False)
            backward_call = generate_command(reverse_p, credentials, api_type=Filesystem_Type, generate_mode=self.generate_mode, openai_model=model)
            return forward_call, backward_call
        raise NotImplementedError


    def run_prompt(self, input_prompt: str, api_type: str):
        credentials = None          # TODO: Work out credentials logic
        api_call, neg_api_call = self.gen_api_pair(input_prompt, api_type, credentials, model="gpt-4-turbo-preview")

        exec_result = self.exec_api_call(api_call, api_type, neg_api_call)


    def test_api_pair_on_docker(self, api_call, neg_api_call, api_type, environment=None) -> bool:
        """Tests an API call and its negation by executing them in a sandbox and verifying the state reversion."""

        if api_type != SQL_Type and api_type != Filesystem_Type:
            # 1. Create Docker sandbox
            sandbox = self.docker_sandbox.create_sandbox()
            # 2. Execute API call and its reversal in the sandbox
            self.api_executor.execute_api_call(api_call)
            self.api_executor.execute_api_call(neg_api_call)
            # 3. Check if the state is reverted successfully
            success = self.check_state_reversion()
            # 4. Log the result in the database
            self.db_manager.log_api_pair(api_call, neg_api_call, success)
            # 5. Delete the sandbox
            self.docker_sandbox.delete_sandbox(sandbox)
        elif api_type == SQL_Type:
            
            test_prompt = "Table Schemas:\n"
            test_prompt += self.db_manager.get_schema_as_string()
            test_prompt += f"\nSQL call: {api_call}\n"
            test_prompt += f"Reverse call: {neg_api_call}\n\n"
            test_prompt += "Given the table schemas, SQL call, and its reversal call, " + \
            "write a Python script to test the SQL call to see if the reversal works. " + \
            "To test, create a new bare minimum database to test this out for " + \
            f"{self.db_manager.db_type}, run the SQL call and its reversal, and check " + \
            "that the start state and the end state are the same. Print out only True " + \
            "if it works, otherwise print out False and print the table contents to stderr. Don't print anything else."

            if self.db_manager.TEST_CONFIG:
                test_prompt += f"Make sure to use the config to access the DB: {self.db_manager.TEST_CONFIG}"


        elif api_type == Filesystem_Type:
            test_prompt = (
                "Path: /sandbox/test_dir"
                f"Shell Command: {api_call}"
                f"Reverse Command: {neg_api_call}"
                "Given a directory path, and a shell command and the reverse command of the shell command, " 
                "can you write a python script to test that the reverse command properly reverses the shell command? "
                "Duplicate the given directory, and perform execution in the duplicate directory only. "
                "Print out only True if it works, otherwise print out False and print errors to stderr. Don't print anything else.\n\n"
            )

            
        res = generate_command(test_prompt, generate_mode=self.generate_mode)

        # get script
        result = self._run_code_in_container(code=res, debug=True, api_type=api_type)

        if not result:
            return False
        
        out, debug = result['output'], result['debug']
        if len(out) >= 1:
            revert_success = bool(out[-1])
            return revert_success
        else:
            print(f"An issue occurred during the dry run: {debug}")
            return False

    def exec_api_call(self, api_call: str, api_type, debug_neg: str = None) -> str:

        if debug_neg:
            neg_api_call = debug_neg


        if self.dry_run_dict[api_type] and debug_neg:
            revert_success = self.test_api_pair_on_docker(api_call, neg_api_call, api_type)
            if not revert_success:
                raise RuntimeError("Dry Run Failed")
            
        if api_type == RESTful_Type:
            self._exec_restful_call(api_call)
            
            payload = (api_call, neg_api_call)

        else:
            payload = (api_call, debug_neg)
            if api_type == SQL_Type:
                if self.db_manager:
                    self._exec_db_call(api_call)
                else:
                    payload = None
            elif api_type == Filesystem_Type:
                self._exec_filesystem_call(api_call)
            
        if payload:
            self._add_api_reverse_to_queue(api_type, payload)

    def undo_api_call(self, api_type: Union[SQL_Type, Filesystem_Type, RESTful_Type], services=None, option=2):
        api_call, neg_api_call = self._pop_api_reverse_from_queue(api_type)

        if not api_call:
            print('History is empty')
            return

        if option == 2:
            # Not a RESTful API negation pair, transaction based
            self._undo_transaction(api_type)

        else:
            if api_type == RESTful_Type:
                self.api_executor.execute_api_call()


    def commit_api_call(self, api_type, arg=None):
        if api_type != RESTful_Type:
            self._commit_transaction(api_type, message=arg)
        else:
            raise NotImplementedError
        self._reset_api_history_queue(api_type)

    def _undo_transaction(self, api_type):
        if api_type == SQL_Type:
            self.db_manager.rollback_db_calls()
        elif api_type == Filesystem_Type:
            self.fs_manager.revert()
        else:
            raise NotImplementedError
        
    def _commit_transaction(self, api_type, message=None):
        if api_type == SQL_Type:
            self.db_manager.commit_db_calls()
        elif api_type == Filesystem_Type:
            if not message:
                message = 'Auto-commit via FSManager'
            self.fs_manager.commit(message=message)
                
        else:
            raise NotImplementedError
        
    def initialize_db(self, debug_manager: DBManager=None):
        print('Initialized DB Manager')
        self.db_manager = debug_manager
        self.db_initialized = True

    def initialize_fs(self, debug_path=None, git_init=True):
        print('Initialized FS Manager')
        self.fs_manager = FSManager(debug_path, git_init=git_init)
        self.fs_manager.initialize_version_control()
        self.fs_initialized = True

    def set_dry_run(self, api_type, on):
        self.dry_run_dict[api_type] = on

    def _exec_db_call(self, call) -> bool:
        """Execute a SQL API call"""
        if not self.db_initialized:
            self.initialize_db()
            
        if re.match(
            r"^SELECT\b", call, re.IGNORECASE
        ):
            return self.db_manager.fetch_db_call(call)
        else:
            self.db_manager.execute_db_call(call)

    def _exec_filesystem_call(self, call) -> bool:
        """Execute a file system API call"""
        if not self.fs_initialized:
            self.initialize_fs()
            self.fs_manager.initialize_version_control()  # make sure that git exists on every call
        self.fs_manager.execute(call, display=True)

    def _exec_restful_call(self, call) -> bool:
        """Execute a RESTful API call"""
        raise NotImplementedError

    def _add_api_reverse_to_queue(self, api_type: Union[SQL_Type, Filesystem_Type, RESTful_Type], payload):
        self.api_history_pair_queue_dict[api_type].append(payload)

    def _pop_api_reverse_from_queue(self, api_type: Union[SQL_Type, Filesystem_Type, RESTful_Type]):
        """Remove a pair of API calls from the list of API pairs"""
        # Dequeue in deque
        return self.api_history_pair_queue_dict[api_type].pop() if self.api_history_pair_queue_dict[api_type] else None

    def _reset_api_history_queue(self, api_type: Union[SQL_Type, Filesystem_Type, RESTful_Type]):
        """Reset the API history queue"""
        # Clear the deque
        self.api_history_pair_queue_dict[api_type].clear()

    def _run_code_in_container(self, code, debug=False, api_type=RESTful_Type):
        try:
            volume_path = None
            if api_type == Filesystem_Type:
                volume_path = os.path.abspath(self.fs_manager.fs_path)
            image = self.docker_sandbox.create_image_from_code(code, api_type)
            response = self.docker_sandbox.create_python_sandbox(code, image, attached_volume=volume_path)
            
            if not debug:
                return response['output']
            else:
                return response
        except Exception as e:
            print("An error occured while trying to execute output locally {e}.\nAborted...".format(e=e))
            return