wehub-resource-sync
279 lines
12 KiB
Python
279 lines
12 KiB
Python
"""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
|